EdgePredMolecularEncoder¶

Predict whether pairs of atoms are connected.

from torch_molecule import EdgePredMolecularEncoder

Examples · Methods

Constructor¶

class EdgePredMolecularEncoder(*, 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='EdgePredMolecularEncoder')[source]¶

This encoder implements a GNN-based model for molecular representation learning using the edge prediction.

References

Paper: Strategies for Pre-training Graph Neural Networks (ICLR 2020) https://arxiv.org/abs/1905.12265
Code: https://github.com/snap-stanford/pretrain-gnns/tree/master/chem

Parameters:

num_layer (int, default=5) – Number of GNN layers.
hidden_size (int, default=300) – Dimension of hidden node features.
drop_ratio (float, default=0.5) – Dropout probability.
norm_layer (str, default="batch_norm") – Type of normalization layer to use. One of [“batch_norm”, “layer_norm”, “instance_norm”, “graph_norm”, “size_norm”, “pair_norm”].
encoder_type (str, default="gin-virtual") – Type of GNN architecture to use.
readout (str, default="sum") – Method for aggregating node features to obtain graph-level representations.
batch_size (int, default=128) – Number of samples per batch for training.
epochs (int, default=500) – Maximum number of training epochs.
learning_rate (float, default=0.001) – Learning rate for optimizer.
grad_clip_value (float, optional) – Maximum norm of gradients for gradient clipping.
weight_decay (float, default=0.0) – L2 regularization strength.
use_lr_scheduler (bool, default=False) – Whether to use learning rate scheduler.
scheduler_factor (float, default=0.5) – Factor by which to reduce learning rate when plateau is reached.
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="EdgePredMolecularEncoder") – Name of the encoder model.

Example¶

examples/encoder/edgepred.py¶

import numpy as np
from torch_molecule import EdgePredMolecularEncoder

def test_edgepred_encoder():
    # Test molecules (simple examples)
    molecules = [
        "CC(=O)O",  # Acetic acid
        "CCO",      # Ethanol
        "CCCC",     # Butane
        "c1ccccc1", # Benzene
        "CCN",      # Ethylamine
    ]

    # Basic initilization test
    model = EdgePredMolecularEncoder(
        num_layer=3,
        hidden_size=300,
        batch_size=5,
        epochs=5,  # Small number for testing
        verbose="progress_bar"
    )
    print("Model initialized successfully")
    
    # Basic self-supervised fitting test
    print("\n=== Testing EdgePred model self-supervised fitting ===")
    model.fit(molecules[:4])
    
    # Model saving and loading test
    print("\n=== Testing model saving and loading ===")
    save_path = "test_model.pt"
    model.save_to_local(save_path)
    print(f"Model saved to {save_path}")

    new_model = EdgePredMolecularEncoder()
    new_model.load_from_local(save_path)
    print("Model loaded successfully")
    
    # Clean up
    import os
    if os.path.exists(save_path):
        os.remove(save_path)
        print(f"Cleaned up {save_path}")

if __name__ == "__main__":
    test_edgepred_encoder()

Methods¶

`fit`	Fit the model to the training data with optional validation set.
`encode`	Encode molecules into vector representations.

For shared checkpoint and parameter operations, see Base Modules.

fit¶

EdgePredMolecularEncoder.fit(X_train)[source]¶

Fit the model to the training data with optional validation set.

Parameters:: X_train (List[str]) – Training set input molecular structures as SMILES strings
Returns:: self – Fitted estimator
Return type:: EdgePredMolecularEncoder

encode¶

EdgePredMolecularEncoder.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