RPGNNMolecularPredictor¶
Relational pooling over node permutations.
from torch_molecule import RPGNNMolecularPredictor
Constructor¶
- class RPGNNMolecularPredictor(num_perm=3, fixed_size=10, num_node_feature=9, num_task=1, task_type='regression', num_layer=5, hidden_size=300, gnn_type='gin-virtual', drop_ratio=0.5, norm_layer='batch_norm', graph_pooling='sum', augmented_feature=None, batch_size=128, epochs=500, learning_rate=0.001, weight_decay=0.0, grad_clip_value=None, patience=50, use_lr_scheduler=False, scheduler_factor=0.5, scheduler_patience=5, loss_criterion=None, evaluate_criterion=None, evaluate_higher_better=None, verbose='none', device=None, model_name='RPGNNMolecularPredictor')[source]¶
This predictor implements a GNN model based on Relational pooling.
The full name of RPGNN is Relational Pooling for Graph Representations.
References
Relational Pooling for Graph Representations. https://arxiv.org/abs/1903.02541
Reference Code: https://github.com/PurdueMINDS/RelationalPooling/tree/master?tab=readme-ov-file
- Parameters:
num_task (int, default=1) – Number of prediction tasks.
task_type (str, default="regression") – Type of prediction task, either “regression” or “classification”.
num_layer (int, default=5) – Number of GNN layers.
hidden_size (int, default=300) – Dimension of hidden node features.
gnn_type (str, default="gin-virtual") – Type of GNN architecture to use. One of [“gin-virtual”, “gcn-virtual”, “gin”, “gcn”].
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”].
graph_pooling (str, default="sum") – Method for aggregating node features to graph-level representations. One of [“sum”, “mean”, “max”].
augmented_feature (list or None, default=None) – Additional molecular fingerprints to use as features. It will be concatenated with the graph representation after pooling. Examples like [“morgan”, “maccs”] or None.
batch_size (int, default=128) – Number of samples per batch for training.
epochs (int, default=500) – Maximum number of training epochs.
loss_criterion (callable, optional) – Loss function for training.
evaluate_criterion (str or callable, optional) – Metric for model evaluation.
evaluate_higher_better (bool, optional) – Whether higher values of the evaluation metric are better.
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.
patience (int, default=50) – Number of epochs to wait for improvement before early stopping.
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 (torch.device or str, optional) – Device to run the model on. If None, will auto-detect GPU or use CPU.
model_name (str, default="RPGNNMolecularPredictor") – Name identifier for the model.
num_perm (int, default=3) – Number of random permutations to use for relational pooling.
fixed_size (int, default=10) – Maximum number of nodes to consider in the graph.
num_node_feature (int, default=9) – Dimension of the input node features. This should match the number of atomic features used to represent each node in the molecular graph (e.g., atomic number, degree, hybridization, etc.).
Example¶
import numpy as np
import torch
from torch_molecule import RPGNNMolecularPredictor
from torch_molecule.utils.search import ParameterType, ParameterSpec
import os
def train_rpgnn_predictor():
smiles_list = [
'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'
]
properties = np.array([0, 0, 1, 1]) # Binary classification
unlabeled_smiles = [
'CC1=CC=C(C=C1)C2=CC(=NN2C3=CC=C(C=C3)S(=O)(=O)N)C(F)(F)F',
'CNC[C@H]1OCc2cnnn2CCCC(=O)N([C@H](C)CO)C[C@@H]1C'
]
num_node_feature = 9
print("\n=== rpgnn model initialization ===")
model = RPGNNMolecularPredictor(
num_task=1,
task_type="regression",
num_layer=3,
num_node_feature=num_node_feature,
fixed_size=10,
hidden_size=300,
batch_size=128,
epochs=2,
verbose="progress_bar"
)
print("RPGNN model initialized successfully")
# 2. Basic fitting test
print("\n=== Testing RPGNN model fitting ===")
model.fit(smiles_list[:3], properties[:3])
print("RPGNN model fitting completed")
# 3. Prediction test
print("\n=== Testing RPGNN model prediction ===")
predictions = model.predict(smiles_list[3:])
print(f"Prediction shape: {predictions['prediction'].shape}")
print(f"Prediction for new molecule: {predictions['prediction']}")
# 4. Auto-fitting test with search space parameters
print("\n=== Testing RPGNN model auto-fitting ===")
# Test auto-fitting
search_parameters = {
'num_layer': ParameterSpec(
param_type=ParameterType.INTEGER,
value_range=(2, 4)
),
'hidden_size': ParameterSpec(
param_type=ParameterType.INTEGER,
value_range=(64, 256)
),
'learning_rate': ParameterSpec(
param_type=ParameterType.LOG_FLOAT,
value_range=(1e-4, 1e-2)
)
}
model_auto = RPGNNMolecularPredictor(
num_task=1,
task_type="regression",
epochs=3,
num_node_feature=num_node_feature,
fixed_size=10,
verbose="progress_bar"
)
model_auto.autofit(
smiles_list,
properties,
X_unlbl=unlabeled_smiles,
search_parameters=search_parameters,
n_trials=2
)
# Test save/load
save_path = "test_rpgnn_model.pt"
model.save(save_path)
new_model = RPGNNMolecularPredictor(
num_task=1,
task_type="regression"
)
new_model.load(save_path)
if os.path.exists(save_path):
os.remove(save_path)
if __name__ == "__main__":
train_rpgnn_predictor()
Methods¶
Fit the model to the training data with optional validation set. |
|
Automatically find the best hyperparameters using Optuna optimization. |
|
Make predictions using the fitted model. |
For shared checkpoint and parameter operations, see Base Modules.
fit¶
- RPGNNMolecularPredictor.fit(X_train, y_train, X_val=None, y_val=None, X_unlbl=None)¶
Fit the model to the training data with optional validation set.
- Parameters:
X_train (List[str]) – Training set input molecular structures as SMILES strings
y_train (Union[List, np.ndarray]) – Training set target values for property prediction
X_val (List[str], optional) – Validation set input molecular structures as SMILES strings. If None, training data will be used for validation
y_val (Union[List, np.ndarray], optional) – Validation set target values. Required if X_val is provided
X_unlbl (List[str], optional) – Unlabeled set input molecular structures as SMILES strings.
- Returns:
self – Fitted estimator
- Return type:
autofit¶
- RPGNNMolecularPredictor.autofit(X_train, y_train, X_val=None, y_val=None, X_unlbl=None, search_parameters=None, n_trials=10)¶
Automatically find the best hyperparameters using Optuna optimization.
- Parameters:
X_train (List[str])
y_train (List | ndarray | None)
X_val (List[str] | None)
y_val (List | ndarray | None)
X_unlbl (List[str] | None)
search_parameters (Dict[str, ParameterSpec] | None)
n_trials (int)
- Return type:
predict¶
- RPGNNMolecularPredictor.predict(X)¶
Make predictions using the fitted model.
- Parameters:
X (List[str]) – List of SMILES strings to make predictions for
- Returns:
- Dictionary containing:
’prediction’: Model predictions (shape: [n_samples, n_tasks])
- Return type:
Dict[str, np.ndarray]