LSTMMolecularPredictor¶
Recurrent sequence learning from SMILES.
from torch_molecule import LSTMMolecularPredictor
Constructor¶
- class LSTMMolecularPredictor(device=None, model_name='LSTMMolecularPredictor', num_task=1, task_type='regression', input_dim=54, output_dim=15, LSTMunits=60, max_input_len=200, batch_size=128, epochs=500, loss_criterion=None, evaluate_criterion=None, evaluate_higher_better=None, learning_rate=0.001, weight_decay=0.0, patience=50, use_lr_scheduler=False, scheduler_factor=0.5, scheduler_patience=5, verbose='none')[source]¶
This predictor implements a LSTM model for molecular property prediction tasks.
References
Predicting Polymers’ Glass Transition Temperature by a Chemical Language Processing Model. https://www.semanticscholar.org/reader/f43ed533b2520567be2d8c24f6396f4e63e96430
- Parameters:
num_task (int, default=1) – Number of prediction tasks.
task_type (str, default="regression") – Type of prediction task, either “regression” or “classification”.
input_dim (int, default=54) – Size of vocabulary for SMILES tokenization.
output_dim (int, default=15) – Dimension of embedding vectors.
LSTMunits (int, default=60) – Number of hidden units in LSTM layers.
max_input_len (int, default=200) – Maximum length of input sequences. Longer sequences will be truncated.
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. Defaults to MSELoss for regression.
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.
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="LSTMMolecularPredictor") – Name identifier for the model.
Example¶
import numpy as np
from torch_molecule import LSTMMolecularPredictor
from torch_molecule.utils.search import ParameterType, ParameterSpec
import torch
def test_lstm_predictor():
# Test data
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' # Additional molecule
]
# Multitask regression with 5 tasks - each molecule has 5 property values
properties = np.array([
[0.1, 0.2, 0.3, 0.4, 0.5], # Properties for molecule 1
[0.2, 0.3, 0.4, 0.5, 0.6], # Properties for molecule 2
[0.8, 0.9, 1.0, 1.1, 1.2], # Properties for molecule 3
[0.9, 1.0, 1.1, 1.2, 1.3] # Properties for molecule 4
])
# 1. Basic initialization test
print("\n=== Testing model initialization ===")
model = LSTMMolecularPredictor(
num_task=5, # 5 tasks for multitask regression
task_type="regression",
output_dim=5, # Output dimension matches number of tasks
LSTMunits=60,
batch_size=2,
epochs=200,
patience=200,
device="cpu",
verbose="print_statement"
)
print("Model initialized successfully")
# 2. Basic fitting test
print("\n=== Testing model fitting ===")
model.fit(smiles_list[:3], properties[:3])
print("Model fitting completed")
# 3. Prediction test
print("\n=== Testing 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 custom parameters
print("\n=== Testing model auto-fitting ===")
search_parameters = {
"output_dim": ParameterSpec(ParameterType.INTEGER, (5, 5)), # Fixed to 5 for 5 tasks
"LSTMunits": ParameterSpec(ParameterType.INTEGER, (30, 120)),
# Float-valued parameters with log scale
"learning_rate": ParameterSpec(ParameterType.LOG_FLOAT, (1e-4, 1e-2)),
}
model_auto = LSTMMolecularPredictor(
num_task=5, # 5 tasks for multitask regression
task_type="regression",
epochs=3, # Small number for testing
verbose="print_statement"
)
model_auto.autofit(
smiles_list,
properties,
search_parameters=search_parameters,
n_trials=2 # Small number for testing
)
print("Model auto-fitting completed")
# 5. 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}")
print(f"Device: {torch.device('cuda' if torch.cuda.is_available() else 'cpu')}")
new_model = LSTMMolecularPredictor(
num_task=5, # 5 tasks for multitask regression
task_type="regression",
output_dim=5, # Output dimension matches number of tasks
LSTMunits=60,
batch_size=2,
epochs=2,
device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
)
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_lstm_predictor()
Methods¶
Fit the model to the training data with optional validation set. |
|
Make predictions using the fitted model. |
For shared checkpoint and parameter operations, see Base Modules.
fit¶
- LSTMMolecularPredictor.fit(X_train, y_train, X_val=None, y_val=None)[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
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
- Return type:
autofit¶
- LSTMMolecularPredictor.autofit(X_train, y_train, X_val=None, y_val=None, search_parameters=None, n_trials=10)[source]¶
- Parameters:
X_train (List[str])
y_train (List | ndarray | None)
X_val (List[str] | None)
y_val (List | ndarray | None)
search_parameters (Dict[str, ParameterSpec] | None)
n_trials (int)
- Return type: