kernel-elastic-autoencoder 3.2.0__tar.gz → 3.2.2__tar.gz

{kernel_elastic_autoencoder-3.2.0 → kernel_elastic_autoencoder-3.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernel_elastic_autoencoder
-Version: 3.2.0
+Version: 3.2.2
 Summary: Implementation of Kernel-Elastic Autoencoder for Molecular Design (https://doi.org/10.1093/pnasnexus/pgae168)
 License: MIT
 Author: Felix Rotter-McCartney

{kernel_elastic_autoencoder-3.2.0 → kernel_elastic_autoencoder-3.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernel_elastic_autoencoder"
-version = "3.2.0"
+version = "3.2.2"
 description = "Implementation of Kernel-Elastic Autoencoder for Molecular Design (https://doi.org/10.1093/pnasnexus/pgae168)"
 authors = [
     { name = "Felix Rotter-McCartney", email = "felix.rotter@mail.utoronto.ca" }

{kernel_elastic_autoencoder-3.2.0 → kernel_elastic_autoencoder-3.2.2}/src/kernel_elastic_autoencoder/training.py

@@ -3,7 +3,7 @@ from collections.abc import Iterable
 
 import torch
 from accelerate import Accelerator
-from accelerate.utils import tqdm
+from accelerate.utils import DistributedDataParallelKwargs, tqdm
 
 from kernel_elastic_autoencoder.config import TrainingConfig
 from kernel_elastic_autoencoder.losses import Loss
@@ -50,7 +50,9 @@ class Trainer:
             train_split: Fraction of dataset used for training. Must be between 0 and 1.
             checkpoint: Path of local checkpoint to be saved and/or resumed.
         """
-        accelerator = Accelerator()
+        accelerator = Accelerator(
+            kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=True)]
+        )
 
         loss_fn = Loss(
             hp_lambda=self.config_typed.hyperparameters.hp_lambda,
@@ -115,9 +117,11 @@ class Trainer:
 
         for epoch in range(curr_epoch, self.config_typed.common.max_epochs):
             model.train()
-            for input_ids, conditions, token_mask, condition_mask in tqdm(
-                dataloader_train, desc=f"Epoch {epoch}, Train Batch"
-            ):
+            if accelerator.is_local_main_process:
+                batch_bar = tqdm(
+                    total=len(dataloader_train), desc=f"Epoch {epoch}, Train Batch"
+                )
+            for input_ids, conditions, token_mask, condition_mask in dataloader_train:
                 optimizer.zero_grad()
                 prediction, prediction_noise, latents_noise = model(
                     input_ids, conditions, token_mask, condition_mask
@@ -127,9 +131,15 @@ class Trainer:
                 )
                 accelerator.backward(loss)
                 optimizer.step()
-            print(f"Train loss: {float(loss.detach())}")
-            
+                if accelerator.is_local_main_process:
+                    batch_bar.update(1)
+            accelerator.print(f"Train loss: {float(loss.detach())}")
+
             model.eval()
+            if accelerator.is_local_main_process:
+                batch_bar = tqdm(
+                    total=len(dataloader_train), desc=f"Epoch {epoch}, Test Batch"
+                )
             for input_ids, conditions, token_mask, condition_mask in tqdm(
                 dataloader_test, desc=f"Epoch {epoch}, Test Batch"
             ):
@@ -143,15 +153,20 @@ class Trainer:
                     loss = loss_fn(
                         prediction, prediction_noise, input_ids[:, 1:], latents_noise
                     )
-            print(f"Test loss: {float(loss.detach())}")
+                    if accelerator.is_local_main_process:
+                        batch_bar.update(1)
+            accelerator.print(f"Test loss: {float(loss.detach())}")
 
             scheduler.step()
 
-            accelerator.wait_for_everyone()
             curr_epoch += 1
-            os.makedirs(os.path.join(checkpoint, "dist/"), exist_ok=True)
-            model.save_pretrained(os.path.join(checkpoint, "dist/"))
-            self.config_typed.to_json(
-                os.path.join(checkpoint, "dist/train_config.json")
-            )
+            accelerator.wait_for_everyone()
             accelerator.save_state(checkpoint)
+            if accelerator.is_main_process:
+                os.makedirs(os.path.join(checkpoint, "dist/"), exist_ok=True)
+                accelerator.unwrap_model(model).save_pretrained(
+                    os.path.join(checkpoint, "dist/")
+                )
+                self.config_typed.to_json(
+                    os.path.join(checkpoint, "dist/train_config.json")
+                )