kernel-elastic-autoencoder 3.2.3__tar.gz → 3.2.5__tar.gz

{kernel_elastic_autoencoder-3.2.3 → kernel_elastic_autoencoder-3.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernel_elastic_autoencoder
-Version: 3.2.3
+Version: 3.2.5
 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.3 → kernel_elastic_autoencoder-3.2.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernel_elastic_autoencoder"
-version = "3.2.3"
+version = "3.2.5"
 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.3 → kernel_elastic_autoencoder-3.2.5}/src/kernel_elastic_autoencoder/config.py

@@ -197,7 +197,7 @@ class TrainingHyperparameterConfig(Config):
         r"vanilla-AE and VAE objectives in the reconstruction loss.",
     )
     hp_sigma: float = Field(
-        default=math.sqrt(32),
+        default=math.sqrt(0.32),
         description=r"Hyperparameter $\sigma$, as used in the Kernel function applied in m-MMD loss. Roughly, "
         r"used as a scaling factor to control the sizes of gradients produced by the m-MMD loss.",
     )

{kernel_elastic_autoencoder-3.2.3 → kernel_elastic_autoencoder-3.2.5}/src/kernel_elastic_autoencoder/training.py

@@ -86,20 +86,25 @@ class Trainer:
         dataset = torch.utils.data.TensorDataset(
             input_ids, conditions, token_mask, condition_mask
         )
+        gen = torch.Generator().manual_seed(0)
         dataset_train, dataset_test = torch.utils.data.random_split(
-            dataset, [train_split, 1 - train_split]
+            dataset, [train_split, 1 - train_split], generator=gen
         )
         dataloader_train = torch.utils.data.DataLoader(
             dataset_train,
             batch_size=self.config_typed.common.batch_size,
+            pin_memory=True,
+            num_workers=4,
         )
         dataloader_test = torch.utils.data.DataLoader(
             dataset_test,
             batch_size=self.config_typed.common.batch_size,
+            pin_memory=True,
+            num_workers=4,
         )
         curr_epoch = 0
 
-        model, optimizer, dataloader_train, dataloader_test, scheduler, curr_epoch = (
+        model, optimizer, dataloader_train, dataloader_test, scheduler, curr_epoch, loss_fn = (
             accelerator.prepare(
                 model,
                 optimizer,
@@ -107,6 +112,7 @@ class Trainer:
                 dataloader_test,
                 scheduler,
                 curr_epoch,
+                loss_fn,
             )
         )
         accelerator.register_for_checkpointing(scheduler)
@@ -118,11 +124,9 @@ class Trainer:
         accelerator.wait_for_everyone()
         for epoch in range(curr_epoch, self.config_typed.common.max_epochs):
             model.train()
-            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:
+            for input_ids, conditions, token_mask, condition_mask in tqdm(
+                dataloader_train, desc=f"Epoch {epoch}, Train Batch"
+            ):
                 optimizer.zero_grad()
                 prediction, prediction_noise, latents_noise = model(
                     input_ids, conditions, token_mask, condition_mask
@@ -132,16 +136,12 @@ class Trainer:
                 )
                 accelerator.backward(loss)
                 optimizer.step()
-                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 dataloader_test:
+            for input_ids, conditions, token_mask, condition_mask in tqdm(
+                dataloader_test, desc=f"Epoch {epoch}, Test Batch"
+            ):
                 with torch.no_grad():
                     prediction, prediction_noise, latents_noise = model(
                         input_ids,
@@ -152,8 +152,6 @@ class Trainer:
                     loss = loss_fn(
                         prediction, prediction_noise, input_ids[:, 1:], latents_noise
                     )
-                    if accelerator.is_local_main_process:
-                        batch_bar.update(1)
             accelerator.print(f"Test loss: {float(loss.detach())}")
 
             scheduler.step()