ContextPredMolecularEncoder¶

Predict molecular subgraph contexts.

from torch_molecule import ContextPredMolecularEncoder

Examples · Methods

Constructor¶

class ContextPredMolecularEncoder(*, mode='cbow', context_size=2, neg_samples=1, num_layer=3, 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='ContextPredMolecularEncoder')[source]¶

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

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:

mode (str, default="cbow") – Type of context prediction task. One of [“cbow”, “skipgram”].
context_size (int, default=2) – Size of the context window used for predicting node-level features.
neg_samples (int, default=1) – Number of negative samples used in the training objective.
num_layer (int, default=3) – 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. One of [“gin-virtual”, “gcn-virtual”, “gin”, “gcn”].
readout (str, default="sum") – Method for aggregating node features to obtain graph-level representations. One of [“sum”, “mean”, “max”].
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 the 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 a learning rate scheduler during training.
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="ContextPredMolecularEncoder") – Name of the encoder model.

Example¶

examples/encoder/contextpred.py¶

import numpy as np
from torch_molecule import ContextPredMolecularEncoder

def test_contextpred_encoder():
    # Test molecules (simple examples)
    molecules = [
        'CNC[C@H]1OCc2cnnn2CCCC(=O)N([C@H](C)CO)C[C@@H]1C',
        'CNC[C@@H]1OCc2cnnn2CCCC(=O)N([C@H](C)CO)C[C@H]1C',
        'C[C@H]1CN([C@@H](C)CO)C(=O)CCCn2cc(nn2)CO[C@@H]1CN(C)C(=O)CCC(F)(F)F',
        'CC1=CC=C(C=C1)C2=CC(=NN2C3=CC=C(C=C3)S(=O)(=O)N)C(F)(F)F'
    ]

    # Basic initilization test
    model = ContextPredMolecularEncoder(
        num_layer=3,
        hidden_size=300,
        batch_size=5,
        epochs=5,  # Small number for testing
        mode='cbow', # 'cbow' or 'skipgram'
        context_size=2,
        neg_samples=1,
        verbose='print_statement'
    )
    print("Model initialized successfully")
    
    # Basic self-supervised fitting test
    print("\n=== Testing ContextPred 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 = ContextPredMolecularEncoder()
    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_contextpred_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¶

ContextPredMolecularEncoder.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:: ContextPredMolecularEncoder

encode¶

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