GraphMAEMolecularEncoder¶
Graph masked autoencoding.
from torch_molecule import GraphMAEMolecularEncoder
Constructor¶
- class GraphMAEMolecularEncoder(*, mask_rate=0.15, mask_edge=False, predictor_type='gin', num_layer=5, hidden_size=300, drop_ratio=0.5, norm_layer='batch_norm', encoder_type='gin-virtual', readout='sum', batch_size=128, epochs=500, learning_rate=0.001, grad_clip_value=None, weight_decay=0.0, use_lr_scheduler=False, scheduler_factor=0.5, scheduler_patience=5, verbose='none', device=None, model_name='GraphMAEMolecularEncoder')[source]¶
GraphMAE: Self-Supervised Masked Graph Autoencoders
References
- Parameters:
mask_rate (float, default=0.15) – Fraction of nodes to mask during training.
mask_edge (bool, default=False) – Whether to mask edges in addition to nodes.
predictor_type (str, default="gin") – Type of predictor network to use for reconstruction. Options: [“gin”, “gcn”, “linear”]
num_layer (int, default=5) – Number of message passing layers in the GNN.
hidden_size (int, default=300) – Dimension of hidden node representations.
drop_ratio (float, default=0.5) – Dropout probability.
norm_layer (str, default="batch_norm") – Type of normalization to use. Options: [“batch_norm”, “layer_norm”, “instance_norm”, “graph_norm”, “size_norm”, “pair_norm”]
encoder_type (str, default="gin-virtual") – Type of GNN encoder to use. Options: [“gin-virtual”, “gcn-virtual”, “gin”, “gcn”]
readout (str, default="sum") – Pooling method to use for graph-level representations. Options: [“sum”, “mean”, “max”]
batch_size (int, default=128) – Batch size for training and inference.
epochs (int, default=500) – Number of training epochs.
learning_rate (float, default=0.001) – Learning rate for optimizer.
grad_clip_value (Optional[float], default=None) – Maximum norm of gradients for gradient clipping. No clipping if None.
weight_decay (float, default=0.0) – L2 regularization factor.
use_lr_scheduler (bool, default=False) – Whether to use a learning rate scheduler.
scheduler_factor (float, default=0.5) – Factor by which to reduce learning rate when using scheduler.
scheduler_patience (int, default=5) – Number of epochs with no improvement after which learning rate will be reduced.
verbose (str, default="none") – Whether to display progress info. Options are: “none”, “progress_bar”, “print_statement”. If any other, “none” is automatically chosen.
device (Optional[Union[torch.device, str]], default=None) – Device to run the model on (CPU or GPU).
model_name (str, default="GraphMAEMolecularEncoder") – Name of the model.
Examples
>>> from torch_molecule import GraphMAEMolecularEncoder >>> encoder = GraphMAEMolecularEncoder(hidden_size=128, epochs=100) >>> encoder.fit(["CC(=O)OC1=CC=CC=C1C(=O)O", "CCO", "C1=CC=CC=C1"]) >>> representations = encoder.encode(["CCO"])
Example¶
import numpy as np
import csv
import os
from torch_molecule import GraphMAEMolecularEncoder
EPOCHS = 5
def test_graphmae_encoder():
# Load molecules from CSV file
data_path = "data/molecule100.csv"
if not os.path.exists(data_path):
print(f"Data file not found: {data_path}")
# Use simple molecules as fallback
molecules = [
"CC(=O)O", # Acetic acid
"CCO", # Ethanol
"CCCC", # Butane
"c1ccccc1", # Benzene
"CCN", # Ethylamine
]
else:
molecules = []
with open(data_path, 'r') as file:
csv_reader = csv.DictReader(file)
for i, row in enumerate(csv_reader):
if i >= 50: # Use first 50 molecules
break
molecules.append(row['smiles'])
print(f"Loaded {len(molecules)} molecules from {data_path}")
# Initialize GraphMAE model
model = GraphMAEMolecularEncoder(
num_layer=3,
hidden_size=128,
batch_size=16,
epochs=EPOCHS, # Small number for testing
mask_rate=0.15,
verbose='progress_bar',
# device="cpu"
)
print("GraphMAE model initialized successfully")
# Test fitting
print("\n=== Testing GraphMAE model self-supervised fitting ===")
model.fit(molecules)
# Test encoding
print("\n=== Testing molecule encoding ===")
encodings = model.encode(molecules[:5])
print(f"Encoding shape: {encodings.shape}")
# Test saving and loading
print("\n=== Testing model saving and loading ===")
save_path = "graphmae_model.pt"
model.save_to_local(save_path)
print(f"Model saved to {save_path}")
new_model = GraphMAEMolecularEncoder()
new_model.load_from_local(save_path)
print("Model loaded successfully")
# Test encoding with loaded model
new_encodings = new_model.encode(molecules[:5])
print(f"New encoding shape: {new_encodings.shape}")
# Verify encodings are the same (or very close)
encoding_diff = (encodings - new_encodings).abs().max().item()
print(f"Max difference between encodings: {encoding_diff}")
# Clean up
if os.path.exists(save_path):
os.remove(save_path)
print(f"Cleaned up {save_path}")
def test_graphmae_with_edge_masking():
# Load molecules from CSV file
data_path = "data/molecule100.csv"
if not os.path.exists(data_path):
print(f"Data file not found: {data_path}")
# Use simple molecules as fallback
molecules = [
"CC(=O)O", # Acetic acid
"CCO", # Ethanol
"CCCC", # Butane
"c1ccccc1", # Benzene
"CCN", # Ethylamine
]
else:
molecules = []
with open(data_path, 'r') as file:
csv_reader = csv.DictReader(file)
for i, row in enumerate(csv_reader):
if i >= 50: # Use first 50 molecules
break
molecules.append(row['smiles'])
print(f"Loaded {len(molecules)} molecules from {data_path}")
# Initialize GraphMAE model with edge masking enabled
model = GraphMAEMolecularEncoder(
num_layer=3,
hidden_size=128,
batch_size=16,
epochs=EPOCHS, # Small number for testing
mask_rate=0.15,
mask_edge=True, # Enable edge masking
verbose='progress_bar',
# device="cpu"
)
print("GraphMAE model with edge masking initialized successfully")
# Test fitting
print("\n=== Testing GraphMAE model with edge masking ===")
model.fit(molecules)
# Test encoding
print("\n=== Testing molecule encoding with edge masking model ===")
encodings = model.encode(molecules[:5])
print(f"Encoding shape: {encodings.shape}")
# Test saving and loading
print("\n=== Testing edge masking model saving and loading ===")
save_path = "graphmae_edge_model.pt"
model.save_to_local(save_path)
print(f"Model saved to {save_path}")
new_model = GraphMAEMolecularEncoder()
new_model.load_from_local(save_path)
print("Model loaded successfully")
# Verify edge masking parameter was preserved
print(f"Loaded model mask_edge parameter: {new_model.mask_edge}")
# Clean up
if os.path.exists(save_path):
os.remove(save_path)
print(f"Cleaned up {save_path}")
if __name__ == "__main__":
print("=== Testing GraphMAE Encoder (Default Configuration) ===")
test_graphmae_encoder()
print("\n=== Testing GraphMAE Encoder with Edge Masking ===")
test_graphmae_with_edge_masking()
Methods¶
Fit the model to the training data with optional validation set. |
|
Encode molecules into vector representations. |
For shared checkpoint and parameter operations, see Base Modules.
fit¶
encode¶
- GraphMAEMolecularEncoder.encode(X, return_type='pt')[source]¶
Encode molecules into vector representations.
- Parameters:
X (List[str]) – List of SMILES strings
return_type (Literal["np", "pt"], default="pt") – Return type of the representations
- Returns:
representations – Molecular representations
- Return type:
ndarray or torch.Tensor