Upload train.py

scripts/train.py

@@ -1,24 +1,48 @@
 import pytorch_lightning as L
 from pytorch_lightning.strategies import DDPStrategy
-from configs.config import Config
-from utils.data_loader import get_dataloaders
-from models.diffusion import Diffusion
+from pytorch_lightning.callbacks import ModelCheckpoint
+import config
+from data_loader import get_dataloaders
+from esm_utils import load_esm2_model
+from diffusion import Diffusion
+import sys
 
 # Get dataloaders
-train_loader, val_loader, _ = get_dataloaders(Config)
+train_loader, val_loader, _ = get_dataloaders(config)
 
-# Initialize model
-latent_diffusion_model = Diffusion(Config, latent_dim=Config.latent_dim)
+# Initialize ESM tokenizer and model
+tokenizer, model = load_esm2_model(config.MODEL_NAME)
+
+# Initialize diffusion model
+latent_diffusion_model = Diffusion(config, latent_dim=config.LATENT_DIM, tokenizer=tokenizer)
+print(latent_diffusion_model)
+sys.stdout.flush()
+
+# Define checkpoints to save best model by minimum validation loss
+checkpoint_callback = ModelCheckpoint(
+    monitor='val_loss',
+    save_top_k=1,
+    mode='min',
+    dirpath="/workspace/a03-sgoel/MDpLM/",
+    filename="best_model_epoch{epoch:02d}"
+)
 
 # Initialize trainer
 trainer = L.Trainer(
-    max_epochs=Config.training["epochs"],
-    gpus=Config.training["gpus"],
-    precision=Config.training["precision"],
+    max_epochs=config.Training.NUM_EPOCHS,
+    precision=config.Training.PRECISION,
+    devices=1,
+    accelerator='gpu',
     strategy=DDPStrategy(find_unused_parameters=False),
-    accumulate_grad_batches=Config.training["accumulate_grad_batches"],
-    default_root_dir=Config.training["save_dir"]
+    accumulate_grad_batches=config.Training.ACCUMULATE_GRAD_BATCHES,
+    default_root_dir=config.Training.SAVE_DIR,
+    callbacks=[checkpoint_callback]
 )
 
+print(trainer)
+print("Training model...")
+sys.stdout.flush()
+
 # Train the model
 trainer.fit(latent_diffusion_model, train_loader, val_loader)
+