LSTMMolecularGenerator¶
Recurrent SMILES generation with optional property labels.
from torch_molecule import LSTMMolecularGenerator
Constructor¶
- class LSTMMolecularGenerator(num_task=0, max_len=100, num_layer=3, hidden_size=512, dropout=0.2, batch_size=128, epochs=10000, learning_rate=0.0002, weight_decay=0.0, use_lr_scheduler=False, scheduler_factor=0.5, scheduler_patience=5, grad_norm_clip=None, verbose='none', *, device=None, model_name='LSTMMolecularGenerator')[source]¶
LSTM-based molecular generator.
This generator implements an LSTM architecture for molecular generation. When conditions (y values) are provided during training, they are used to initialize the hidden and cell states of the LSTM.
- Parameters:
num_task (int, default=0) – Number of tasks for conditional generation (0 means unconditional generation).
max_len (int, default=100) – Maximum length of the SMILES strings.
num_layer (int, default=3) – Number of LSTM layers.
hidden_size (int, default=512) – Dimension of hidden states in LSTM.
dropout (float, default=0.2) – Dropout probability for regularization.
batch_size (int, default=128) – Batch size for training.
epochs (int, default=10000) – Maximum number of training epochs.
learning_rate (float, default=0.0002) – Learning rate for optimizer.
weight_decay (float, default=0.0) – L2 regularization factor.
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 using scheduler (if True).
scheduler_patience (int, default=5) – Number of epochs with no improvement after which learning rate will be reduced (if True).
grad_norm_clip (Optional[float], default=None) – Maximum norm for gradient clipping. None means no clipping.
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="LSTMMolecularGenerator") – Name identifier for the model.
Example¶
import os
import csv
import numpy as np
from tqdm import tqdm
import torch
from torch_molecule.generator.lstm import LSTMMolecularGenerator
# EPOCHS = 1000 # Reduced for faster testing
EPOCHS = 5
BATCH_SIZE = 24
def test_lstm_generator():
# Load data from polymer100.csv
data_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
"data", "polymer100.csv")
print(f"Loading data from: {data_path}")
# Read CSV without pandas
smiles_list = []
properties = []
property_columns = []
with open(data_path, 'r') as csvfile:
reader = csv.DictReader(csvfile)
# Get property columns (all columns except 'smiles')
property_columns = [col for col in reader.fieldnames if col != 'smiles']
for row in reader:
smiles_list.append(row['smiles'])
# Extract property values for this row
row_properties = [float(row[col]) for col in property_columns]
properties.append(row_properties)
print(f"Loaded {len(smiles_list)} molecules with {len(property_columns)} properties")
print(f"Property columns: {property_columns}")
print(f"First 3 SMILES: {smiles_list[:3]}")
print(f"First 3 properties: {properties[:3]}")
# 1. Basic initialization test - Unconditional Model
print("\n=== Testing Unconditional LSTM model initialization ===")
unconditional_model = LSTMMolecularGenerator(
num_layer=3,
hidden_size=128,
max_len=64,
batch_size=BATCH_SIZE,
epochs=EPOCHS,
verbose="progress_bar"
)
print("Unconditional LSTM Model initialized successfully")
# 2. Basic fitting test - Unconditional Model
print("\n=== Testing Unconditional LSTM model fitting ===")
unconditional_model.fit(smiles_list)
print("Unconditional LSTM Model fitting completed")
# 3. Unconditional generation test
print("\n=== Testing Unconditional LSTM generation ===")
generated_smiles_uncond = unconditional_model.generate(batch_size=BATCH_SIZE)
print(f"Unconditionally generated {len(generated_smiles_uncond)} molecules")
print("Example unconditionally generated SMILES:", generated_smiles_uncond[:10])
# 4. Model saving and loading test - Unconditional Model
print("\n=== Testing Unconditional LSTM model saving and loading ===")
save_path = "unconditional_lstm_test_model.pt"
unconditional_model.save_to_local(save_path)
print(f"Unconditional LSTM Model saved to {save_path}")
new_unconditional_model = LSTMMolecularGenerator()
new_unconditional_model.load_from_local(save_path)
print("Unconditional LSTM Model loaded successfully")
# Test generation with loaded unconditional model
generated_smiles_uncond = new_unconditional_model.generate(batch_size=5)
print("Generated molecules with loaded unconditional model:", len(generated_smiles_uncond))
print("Example generated SMILES:", generated_smiles_uncond[:10])
# Clean up unconditional model
if os.path.exists(save_path):
os.remove(save_path)
print(f"Cleaned up {save_path}")
# 5. Basic initialization test - Property Conditional Model
print("\n=== Testing Property Conditional LSTM model initialization ===")
prop_conditional_model = LSTMMolecularGenerator(
num_layer=2,
hidden_size=128,
max_len=64,
num_task=len(property_columns), # Set number of properties
batch_size=BATCH_SIZE,
epochs=EPOCHS,
verbose="progress_bar"
)
print("Property Conditional LSTM Model initialized successfully")
# 6. Basic fitting test - Property Conditional Model
print("\n=== Testing Property Conditional LSTM model fitting ===")
prop_conditional_model.fit(smiles_list, properties)
print("Property Conditional LSTM Model fitting completed")
# 7. Property conditional generation test
print("\n=== Testing Property Conditional LSTM generation ===")
# Create some target properties (using mean values from the dataset as a starting point)
mean_properties = np.mean(properties, axis=0).tolist()
target_properties = []
for i in range(5):
# Create variations around the mean
target_prop = [p * (0.8 + 0.4 * np.random.random()) for p in mean_properties]
target_properties.append(target_prop)
print(f"Target properties for generation: {target_properties}")
generated_smiles = prop_conditional_model.generate(labels=target_properties)
print(f"Property conditionally generated {len(generated_smiles)} molecules")
print("Example property conditionally generated SMILES:", generated_smiles[:2])
# 8. Model saving and loading test - Property Conditional Model
print("\n=== Testing Property Conditional LSTM model saving and loading ===")
save_path = "prop_conditional_lstm_test_model.pt"
prop_conditional_model.save_to_local(save_path)
print(f"Property Conditional LSTM Model saved to {save_path}")
new_prop_conditional_model = LSTMMolecularGenerator()
new_prop_conditional_model.load_from_local(save_path)
print("Property Conditional LSTM Model loaded successfully")
# Test generation with loaded property conditional model
generated_smiles = new_prop_conditional_model.generate(labels=target_properties)
print("Generated molecules with loaded property conditional model:", len(generated_smiles))
print("Example generated SMILES:", generated_smiles[:2])
# Clean up property conditional model
if os.path.exists(save_path):
os.remove(save_path)
print(f"Cleaned up {save_path}")
if __name__ == "__main__":
test_lstm_generator()
Methods¶
For shared checkpoint and parameter operations, see Base Modules.
fit¶
generate¶
- LSTMMolecularGenerator.generate(labels=None, batch_size=32)[source]¶
Generate molecules using LSTM.
- Parameters:
labels (Optional[Union[List[List], np.ndarray]]) – Target property values for conditional generation.
batch_size (int) – Number of molecules to generate.
- Returns:
Generated molecules as SMILES strings.
- Return type:
List[str]