InfoGraphMolecularEncoder¶

Learn agreement between graph and substructure representations.

from torch_molecule import InfoGraphMolecularEncoder

Examples · Methods

Constructor¶

class InfoGraphMolecularEncoder(lw_prior=0.0, embedding_dim=160, num_layer=5, 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, *, device=None, model_name='InfoGraphMolecularEncoder', verbose='none')[source]¶

This encoder implements a InfoGraph for molecular representation learning.

InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization (ICLR 2020)

References

Parameters:

lw_prior (float, default=0.) – Weight for prior loss term.
embedding_dim (int, default=160) – Dimension of final graph embedding. Must be divisible by num_layer.
num_layer (int, default=5) – Number of GNN layers.
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 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.
scheduler_factor (float, default=0.5) – Factor by which to reduce the learning rate when plateau is detected.
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="InfoGraphMolecularEncoder") – Name identifier for the model.

Example¶

examples/encoder/infograph.py¶

import numpy as np
from torch_molecule import InfoGraphMolecularEncoder

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

    # Basic initialization test
    model = InfoGraphMolecularEncoder(
        num_layer=3,
        embedding_dim=300,  # Must be divisible by num_layer
        batch_size=5,
        epochs=5,  # Small number for testing
        verbose='progress_bar',
        lw_prior=0.1
    )
    print("Model initialized successfully")
    
    # Basic self-supervised fitting test
    print("\n=== Testing Infograph model self-supervised fitting ===")
    model.fit(molecules)
    
    # Encoding test
    print("\n=== Testing molecule encoding ===")
    encodings = model.encode(molecules)
    print(f"Encoded {len(molecules)} molecules to shape {encodings.shape}")
    
    # Model saving and loading test
    print("\n=== Testing model saving and loading ===")
    save_path = "test_infograph_model.pt"
    model.save_to_local(save_path)
    print(f"Model saved to {save_path}")

    new_model = InfoGraphMolecularEncoder()
    new_model.load_from_local(save_path)
    print("Model loaded successfully")
    
    # Test encoding with loaded model
    new_encodings = new_model.encode(molecules)
    print(f"Encoded {len(molecules)} molecules with loaded model")
    
    # Verify encodings are the same
    if np.allclose(encodings.numpy(), new_encodings.numpy()):
        print("Encodings from original and loaded models match!")
    else:
        print("Warning: Encodings from original and loaded models differ")
    
    # Clean up
    import os
    if os.path.exists(save_path):
        os.remove(save_path)
        print(f"Cleaned up {save_path}")

if __name__ == "__main__":
    test_infograph_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¶

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

encode¶

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