JTVAEMolecularGenerator¶
Unconditional junction-tree variational autoencoding.
from torch_molecule import JTVAEMolecularGenerator
Constructor¶
- class JTVAEMolecularGenerator(hidden_size=450, latent_size=56, depthT=20, depthG=3, batch_size=32, epochs=20, learning_rate=0.003, weight_decay=0.0, grad_norm_clip=None, beta=0.0, step_beta=0.002, max_beta=1.0, warmup=40000, use_lr_scheduler=True, anneal_rate=0.9, anneal_iter=40000, kl_anneal_iter=2000, verbose='none', *, device=None, model_name='JTVAEMolecularGenerator')[source]¶
JT-VAE-based molecular generator. Implemented for unconditional molecular generation.
References
Junction Tree Variational Autoencoder for Molecular Graph Generation. ICML 2018. https://arxiv.org/pdf/1802.04364
- Parameters:
hidden_size (int, default=450) – Dimension of hidden layers in the model.
latent_size (int, default=56) – Dimension of the latent space.
depthT (int, default=20) – Depth of the tree encoder.
depthG (int, default=3) – Depth of the graph decoder.
batch_size (int, default=32) – Number of samples per batch during training.
epochs (int, default=20) – Number of epochs to train the model.
learning_rate (float, default=0.003) – Initial learning rate for the optimizer.
weight_decay (float, default=0.0) – L2 regularization factor.
grad_norm_clip (Optional[float], default=None) – Maximum norm for gradient clipping. None means no clipping.
beta (float, default=0.0) – Initial KL divergence weight for VAE training.
step_beta (float, default=0.002) – Step size for KL annealing.
max_beta (float, default=1.0) – Maximum value for KL weight.
warmup (int, default=40000) – Number of steps for KL annealing warmup.
use_lr_scheduler (bool, default=True) – Whether to use learning rate scheduling.
anneal_rate (float, default=0.9) – Learning rate annealing factor.
anneal_iter (int, default=40000) – Number of iterations between learning rate updates.
kl_anneal_iter (int, default=2000) – Number of iterations between KL weight updates.
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="JTVAEMolecularGenerator") – Name identifier for the model.
Example¶
import os
import numpy as np
import csv
from tqdm import tqdm
import torch
from torch_molecule.generator.jtvae import JTVAEMolecularGenerator
EPOCHS = 3 # Reduced for faster testing
BATCH_SIZE = 8
def test_jtvae_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 = []
with open(data_path, 'r', encoding='utf-8') as csvfile:
reader = csv.DictReader(csvfile)
for row in reader:
smiles_list.append(row['smiles'])
print(f"Loaded {len(smiles_list)} molecules")
print(f"First 3 SMILES: {smiles_list[:3]}")
# 1. Basic initialization test
print("\n=== Testing JTVAE model initialization ===")
jtvae_model = JTVAEMolecularGenerator(
hidden_size=300, # Reduced for faster testing
latent_size=32, # Reduced for faster testing
depthT=10, # Reduced for faster testing
depthG=2, # Reduced for faster testing
batch_size=BATCH_SIZE,
epochs=EPOCHS,
verbose="print_statement"
)
print("JTVAE Model initialized successfully")
# 2. Basic fitting test
print("\n=== Testing JTVAE model fitting ===")
jtvae_model.fit(smiles_list[:50]) # Using a subset for faster testing
print("JTVAE Model fitting completed")
# 3. Generation test
print("\n=== Testing JTVAE generation ===")
generated_smiles = jtvae_model.generate(batch_size=5)
print(f"Generated {len(generated_smiles)} molecules")
print("Example generated SMILES:", generated_smiles)
# 4. Model saving and loading test
print("\n=== Testing JTVAE model saving and loading ===")
save_path = "jtvae_test_model.pt"
jtvae_model.save_to_local(save_path)
print(f"JTVAE Model saved to {save_path}")
new_jtvae_model = JTVAEMolecularGenerator()
new_jtvae_model.load_from_local(save_path)
print("JTVAE Model loaded successfully")
# Test generation with loaded model
generated_smiles = new_jtvae_model.generate(batch_size=5)
print("Generated molecules with loaded model:", len(generated_smiles))
print("Example generated SMILES:", generated_smiles)
# Clean up
if os.path.exists(save_path):
os.remove(save_path)
print(f"Cleaned up {save_path}")
if __name__ == "__main__":
test_jtvae_generator()
Methods¶
For shared checkpoint and parameter operations, see Base Modules.