fixes #110 & fix some tests & black 23.10

clementchadebec · Oct 22, 2023 · 1103a99 · 1103a99
1 parent b7c53f4
commit 1103a99
Show file tree

Hide file tree

Showing 90 changed files with 462 additions and 1,259 deletions.
diff --git a/src/pythae/data/datasets.py b/src/pythae/data/datasets.py
@@ -55,7 +55,6 @@ class BaseDataset(Dataset):
     """
 
     def __init__(self, data, labels):
-
         self.labels = labels.type(torch.float)
         self.data = data.type(torch.float)
 

diff --git a/src/pythae/data/preprocessors.py b/src/pythae/data/preprocessors.py
@@ -92,7 +92,6 @@ def to_dataset(data: torch.Tensor, labels: Optional[torch.Tensor] = None):
         return dataset
 
     def _process_data_array(self, data: np.ndarray, batch_size: int = 100):
-
         num_samples = data.shape[0]
         samples_shape = data.shape
 
@@ -102,7 +101,6 @@ def _process_data_array(self, data: np.ndarray, batch_size: int = 100):
         full_data = []
 
         for i in range(num_complete_batch):
-
             # Detect potential nan
             if DataProcessor.has_nan(data[i * batch_size : (i + 1) * batch_size]):
                 raise ValueError("Nan detected in input data!")

diff --git a/src/pythae/models/adversarial_ae/adversarial_ae_model.py b/src/pythae/models/adversarial_ae/adversarial_ae_model.py
@@ -58,7 +58,6 @@ def __init__(
         decoder: Optional[BaseDecoder] = None,
         discriminator: Optional[BaseDiscriminator] = None,
     ):
-
         VAE.__init__(self, model_config=model_config, encoder=encoder, decoder=decoder)
 
         if discriminator is None:
@@ -149,22 +148,16 @@ def forward(self, inputs: BaseDataset, **kwargs) -> ModelOutput:
         return output
 
     def loss_function(self, recon_x, x, z, z_prior):
-
         N = z.shape[0]  # batch size
 
         if self.model_config.reconstruction_loss == "mse":
-
-            recon_loss = (
-                0.5
-                * F.mse_loss(
-                    recon_x.reshape(x.shape[0], -1),
-                    x.reshape(x.shape[0], -1),
-                    reduction="none",
-                ).sum(dim=-1)
-            )
+            recon_loss = 0.5 * F.mse_loss(
+                recon_x.reshape(x.shape[0], -1),
+                x.reshape(x.shape[0], -1),
+                reduction="none",
+            ).sum(dim=-1)
 
         elif self.model_config.reconstruction_loss == "bce":
-
             recon_loss = F.binary_cross_entropy(
                 recon_x.reshape(x.shape[0], -1),
                 x.reshape(x.shape[0], -1),
@@ -232,7 +225,6 @@ def save(self, dir_path: str):
 
     @classmethod
     def _load_custom_discriminator_from_folder(cls, dir_path):
-
         file_list = os.listdir(dir_path)
         cls._check_python_version_from_folder(dir_path=dir_path)
 
@@ -361,7 +353,6 @@ def load_from_hf_hub(
             )
 
         else:
-
             if not model_config.uses_default_encoder:
                 _ = hf_hub_download(repo_id=hf_hub_path, filename="encoder.pkl")
                 encoder = cls._load_custom_encoder_from_folder(dir_path)

diff --git a/src/pythae/models/ae/ae_model.py b/src/pythae/models/ae/ae_model.py
@@ -39,7 +39,6 @@ def __init__(
         encoder: Optional[BaseEncoder] = None,
         decoder: Optional[BaseDecoder] = None,
     ):
-
         BaseAE.__init__(self, model_config=model_config, decoder=decoder)
 
         self.model_name = "AE"
@@ -83,7 +82,6 @@ def forward(self, inputs: BaseDataset, **kwargs) -> ModelOutput:
         return output
 
     def loss_function(self, recon_x, x):
-
         MSE = F.mse_loss(
             recon_x.reshape(x.shape[0], -1), x.reshape(x.shape[0], -1), reduction="none"
         ).sum(dim=-1)

diff --git a/src/pythae/models/base/base_model.py b/src/pythae/models/base/base_model.py
@@ -57,10 +57,9 @@ class BaseAE(nn.Module):
     def __init__(
         self,
         model_config: BaseAEConfig,
-        encoder: Optional[BaseDecoder] = None,
+        encoder: Optional[BaseEncoder] = None,
         decoder: Optional[BaseDecoder] = None,
     ):
-
         nn.Module.__init__(self)
 
         self.model_name = "BaseAE"
@@ -374,7 +373,6 @@ def _load_model_weights_from_folder(cls, dir_path):
 
     @classmethod
     def _load_custom_encoder_from_folder(cls, dir_path):
-
         file_list = os.listdir(dir_path)
         cls._check_python_version_from_folder(dir_path=dir_path)
 
@@ -393,7 +391,6 @@ def _load_custom_encoder_from_folder(cls, dir_path):
 
     @classmethod
     def _load_custom_decoder_from_folder(cls, dir_path):
-
         file_list = os.listdir(dir_path)
         cls._check_python_version_from_folder(dir_path=dir_path)
 
@@ -510,7 +507,6 @@ def load_from_hf_hub(cls, hf_hub_path: str, allow_pickle=False):  # pragma: no c
             )
 
         else:
-
             if not model_config.uses_default_encoder:
                 _ = hf_hub_download(repo_id=hf_hub_path, filename="encoder.pkl")
                 encoder = cls._load_custom_encoder_from_folder(dir_path)

diff --git a/src/pythae/models/beta_tc_vae/beta_tc_vae_model.py b/src/pythae/models/beta_tc_vae/beta_tc_vae_model.py
@@ -41,7 +41,6 @@ def __init__(
         encoder: Optional[BaseEncoder] = None,
         decoder: Optional[BaseDecoder] = None,
     ):
-
         VAE.__init__(self, model_config=model_config, encoder=encoder, decoder=decoder)
 
         self.model_name = "BetaTCVAE"
@@ -89,20 +88,14 @@ def forward(self, inputs: BaseDataset, **kwargs):
         return output
 
     def loss_function(self, recon_x, x, mu, log_var, z, dataset_size):
-
         if self.model_config.reconstruction_loss == "mse":
-
-            recon_loss = (
-                0.5
-                * F.mse_loss(
-                    recon_x.reshape(x.shape[0], -1),
-                    x.reshape(x.shape[0], -1),
-                    reduction="none",
-                ).sum(dim=-1)
-            )
+            recon_loss = 0.5 * F.mse_loss(
+                recon_x.reshape(x.shape[0], -1),
+                x.reshape(x.shape[0], -1),
+                reduction="none",
+            ).sum(dim=-1)
 
         elif self.model_config.reconstruction_loss == "bce":
-
             recon_loss = F.binary_cross_entropy(
                 recon_x.reshape(x.shape[0], -1),
                 x.reshape(x.shape[0], -1),

diff --git a/src/pythae/models/beta_vae/beta_vae_model.py b/src/pythae/models/beta_vae/beta_vae_model.py
@@ -40,7 +40,6 @@ def __init__(
         encoder: Optional[BaseEncoder] = None,
         decoder: Optional[BaseDecoder] = None,
     ):
-
         VAE.__init__(self, model_config=model_config, encoder=encoder, decoder=decoder)
 
         self.model_name = "BetaVAE"
@@ -81,20 +80,14 @@ def forward(self, inputs: BaseDataset, **kwargs):
         return output
 
     def loss_function(self, recon_x, x, mu, log_var, z):
-
         if self.model_config.reconstruction_loss == "mse":
-
-            recon_loss = (
-                0.5
-                * F.mse_loss(
-                    recon_x.reshape(x.shape[0], -1),
-                    x.reshape(x.shape[0], -1),
-                    reduction="none",
-                ).sum(dim=-1)
-            )
+            recon_loss = 0.5 * F.mse_loss(
+                recon_x.reshape(x.shape[0], -1),
+                x.reshape(x.shape[0], -1),
+                reduction="none",
+            ).sum(dim=-1)
 
         elif self.model_config.reconstruction_loss == "bce":
-
             recon_loss = F.binary_cross_entropy(
                 recon_x.reshape(x.shape[0], -1),
                 x.reshape(x.shape[0], -1),

diff --git a/src/pythae/models/ciwae/ciwae_model.py b/src/pythae/models/ciwae/ciwae_model.py
@@ -39,7 +39,6 @@ def __init__(
         encoder: Optional[BaseEncoder] = None,
         decoder: Optional[BaseDecoder] = None,
     ):
-
         VAE.__init__(self, model_config=model_config, encoder=encoder, decoder=decoder)
 
         self.model_name = "CIWAE"
@@ -92,22 +91,16 @@ def forward(self, inputs: BaseDataset, **kwargs):
         return output
 
     def loss_function(self, recon_x, x, mu, log_var, z):
-
         if self.model_config.reconstruction_loss == "mse":
-
-            recon_loss = (
-                0.5
-                * F.mse_loss(
-                    recon_x,
-                    x.reshape(recon_x.shape[0], -1)
-                    .unsqueeze(1)
-                    .repeat(1, self.n_samples, 1),
-                    reduction="none",
-                ).sum(dim=-1)
-            )
+            recon_loss = 0.5 * F.mse_loss(
+                recon_x,
+                x.reshape(recon_x.shape[0], -1)
+                .unsqueeze(1)
+                .repeat(1, self.n_samples, 1),
+                reduction="none",
+            ).sum(dim=-1)
 
         elif self.model_config.reconstruction_loss == "bce":
-
             recon_loss = F.binary_cross_entropy(
                 recon_x,
                 x.reshape(recon_x.shape[0], -1)
@@ -117,7 +110,7 @@ def loss_function(self, recon_x, x, mu, log_var, z):
             ).sum(dim=-1)
 
         log_q_z = (-0.5 * (log_var + torch.pow(z - mu, 2) / log_var.exp())).sum(dim=-1)
-        log_p_z = -0.5 * (z ** 2).sum(dim=-1)
+        log_p_z = -0.5 * (z**2).sum(dim=-1)
 
         KLD = -(log_p_z - log_q_z)
 

diff --git a/src/pythae/models/disentangled_beta_vae/disentangled_beta_vae_model.py b/src/pythae/models/disentangled_beta_vae/disentangled_beta_vae_model.py
@@ -40,7 +40,6 @@ def __init__(
         encoder: Optional[BaseEncoder] = None,
         decoder: Optional[BaseDecoder] = None,
     ):
-
         VAE.__init__(self, model_config=model_config, encoder=encoder, decoder=decoder)
 
         assert (
@@ -89,20 +88,14 @@ def forward(self, inputs: BaseDataset, **kwargs):
         return output
 
     def loss_function(self, recon_x, x, mu, log_var, z, epoch):
-
         if self.model_config.reconstruction_loss == "mse":
-
-            recon_loss = (
-                0.5
-                * F.mse_loss(
-                    recon_x.reshape(x.shape[0], -1),
-                    x.reshape(x.shape[0], -1),
-                    reduction="none",
-                ).sum(dim=-1)
-            )
+            recon_loss = 0.5 * F.mse_loss(
+                recon_x.reshape(x.shape[0], -1),
+                x.reshape(x.shape[0], -1),
+                reduction="none",
+            ).sum(dim=-1)
 
         elif self.model_config.reconstruction_loss == "bce":
-
             recon_loss = F.binary_cross_entropy(
                 recon_x.reshape(x.shape[0], -1),
                 x.reshape(x.shape[0], -1),

diff --git a/src/pythae/models/factor_vae/factor_vae_model.py b/src/pythae/models/factor_vae/factor_vae_model.py
@@ -43,7 +43,6 @@ def __init__(
         encoder: Optional[BaseEncoder] = None,
         decoder: Optional[BaseDecoder] = None,
     ):
-
         VAE.__init__(self, model_config=model_config, encoder=encoder, decoder=decoder)
 
         self.discriminator = FactorVAEDiscriminator(latent_dim=model_config.latent_dim)
@@ -132,22 +131,16 @@ def forward(self, inputs: BaseDataset, **kwargs) -> ModelOutput:
         return output
 
     def loss_function(self, recon_x, x, mu, log_var, z, z_bis_permuted):
-
         N = z.shape[0]  # batch size
 
         if self.model_config.reconstruction_loss == "mse":
-
-            recon_loss = (
-                0.5
-                * F.mse_loss(
-                    recon_x.reshape(x.shape[0], -1),
-                    x.reshape(x.shape[0], -1),
-                    reduction="none",
-                ).sum(dim=-1)
-            )
+            recon_loss = 0.5 * F.mse_loss(
+                recon_x.reshape(x.shape[0], -1),
+                x.reshape(x.shape[0], -1),
+                reduction="none",
+            ).sum(dim=-1)
 
         elif self.model_config.reconstruction_loss == "bce":
-
             recon_loss = F.binary_cross_entropy(
                 recon_x.reshape(x.shape[0], -1),
                 x.reshape(x.shape[0], -1),

diff --git a/src/pythae/models/factor_vae/factor_vae_utils.py b/src/pythae/models/factor_vae/factor_vae_utils.py
@@ -4,7 +4,6 @@
 
 class FactorVAEDiscriminator(nn.Module):
     def __init__(self, latent_dim=16, hidden_units=1000) -> None:
-
         nn.Module.__init__(self)
 
         self.layers = nn.Sequential(