AttrMaskMolecularEncoder¶

Reconstruct masked atom attributes.

from torch_molecule import AttrMaskMolecularEncoder

Examples · Methods

Constructor¶

class AttrMaskMolecularEncoder(*, mask_num=0, mask_rate=0.15, 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='AttrMaskMolecularEncoder')[source]¶

This encoder implements a GNN-based model for molecular representation learning using the attribute masking 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:

mask_num (int, default=0) – Number of atom features to mask during pretraining. If set to 0, masking is determined by mask_rate.
mask_rate (float, default=0.15) – Proportion of atoms to mask randomly. Ignored if mask_num is set.
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. 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="AttrMaskMolecularEncoder") – Name of the encoder model.

Example¶

examples/encoder/attrmask.py¶

import numpy as np
from torch_molecule import AttrMaskMolecularEncoder

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

    # Basic initilization test
    model = AttrMaskMolecularEncoder(
        num_layer=3,
        hidden_size=300,
        batch_size=5,
        epochs=5,  # Small number for testing
        verbose='print_statement'
    )
    print("Model initialized successfully")
    
    # Basic self-supervised fitting test
    print("\n=== Testing AttrMask 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 = AttrMaskMolecularEncoder()
    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}")

def test_attrmask_encoder_polymers():
    # Test molecules (simple examples)
    polymers = [
        "*Nc1ccc([C@H](CCC)c2ccc(C3(c4ccc([C@@H](CCC)c5ccc(N*)cc5)cc4)CCC(CCCCC)CC3)cc2)cc1",
        "*Nc1ccc(-c2c(-c3ccc(C)cc3)c(-c3ccc(C)cc3)c(N*)c(-c3ccc(C)cc3)c2-c2ccc(C)cc2)cc1",
        "*CC(*)(C)C(=O)OCCCCCCCCCOc1ccc2cc(C(=O)Oc3ccccc3)ccc2c1"
    ]
    model = AttrMaskMolecularEncoder(
        num_layer=3,
        hidden_size=300,
        batch_size=5,
        epochs=5,  # Small number for testing
        verbose='progress_bar'
    )
    model.fit(polymers)
    vectors = model.encode(polymers)
    print(f"Representation shape: {vectors.shape}")
    print(f"Representation for new molecule: {vectors[0]}")


if __name__ == "__main__":
    test_attrmask_encoder_polymers()
    test_attrmask_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¶

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

encode¶

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