changed nn.Tanh() output layers to nn.Sigmoid()

models/beta_vae.py

@@ -83,7 +83,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
     def encode(self, input: Tensor) -> List[Tensor]:
         """

models/betatc_vae.py

@@ -79,7 +79,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
     def encode(self, input: Tensor) -> List[Tensor]:
         """

models/cat_vae.py

@@ -83,7 +83,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
         self.sampling_dist = torch.distributions.OneHotCategorical(1. / categorical_dim * torch.ones((self.categorical_dim, 1)))
 
     def encode(self, input: Tensor) -> List[Tensor]:

models/cvae.py

@@ -78,7 +78,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
     def encode(self, input: Tensor) -> List[Tensor]:
         """

models/dfcvae.py

@@ -76,7 +76,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
         self.feature_network = vgg19_bn(pretrained=True)
 

models/dip_vae.py

@@ -73,7 +73,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
     def encode(self, input: Tensor) -> List[Tensor]:
         """

models/fvae.py

@@ -73,7 +73,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
         # Discriminator network for the Total Correlation (TC) loss
         self.discriminator = nn.Sequential(nn.Linear(self.latent_dim, 1000),

models/hvae.py

@@ -101,7 +101,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
         # ========================================================================#
         # Pesudo Input for the Vamp-Prior

models/info_vae.py

@@ -83,7 +83,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
     def encode(self, input: Tensor) -> List[Tensor]:
         """

models/iwae.py

@@ -73,7 +73,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
     def encode(self, input: Tensor) -> List[Tensor]:
         """

models/joint_vae.py

@@ -105,7 +105,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
         self.sampling_dist = torch.distributions.OneHotCategorical(1. / categorical_dim * torch.ones((self.categorical_dim, 1)))
 
     def encode(self, input: Tensor) -> List[Tensor]:

models/logcosh_vae.py

@@ -73,7 +73,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
     def encode(self, input: Tensor) -> List[Tensor]:
         """

models/lvae.py

@@ -128,7 +128,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
         hidden_dims.reverse()
 
     def encode(self, input: Tensor) -> List[Tensor]:

models/miwae.py

@@ -76,7 +76,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
     def encode(self, input: Tensor) -> List[Tensor]:
         """

models/mssim_vae.py

@@ -75,7 +75,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
         self.mssim_loss = MSSIM(self.in_channels,
                                 window_size,

models/swae.py

@@ -79,7 +79,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
     def encode(self, input: Tensor) -> Tensor:
         """

models/twostage_vae.py

@@ -70,7 +70,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
         #---------------------- Second VAE ---------------------------#
         encoder2 = []

models/vampvae.py

@@ -73,11 +73,11 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
         self.pseudo_input = torch.eye(self.num_components, requires_grad= False)
         self.embed_pseudo = nn.Sequential(nn.Linear(self.num_components, 12288),
-                                          nn.Hardtanh(0.0, 1.0)) # 3x64x64 = 12288
+                                          nn.HardSigmoid(0.0, 1.0)) # 3x64x64 = 12288
 
     def encode(self, input: Tensor) -> List[Tensor]:
         """

models/vanilla_vae.py

@@ -72,7 +72,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
     def encode(self, input: Tensor) -> List[Tensor]:
         """

models/vq_vae.py

@@ -161,7 +161,7 @@ def __init__(self,
                                    out_channels=3,
                                    kernel_size=4,
                                    stride=2, padding=1),
-                nn.Tanh()))
+                nn.Sigmoid()))
 
         self.decoder = nn.Sequential(*modules)
 

models/wae_mmd.py

@@ -76,7 +76,7 @@ def __init__(self,
                             nn.LeakyReLU(),
                             nn.Conv2d(hidden_dims[-1], out_channels= 3,
                                       kernel_size= 3, padding= 1),
-                            nn.Tanh())
+                            nn.Sigmoid())
 
     def encode(self, input: Tensor) -> Tensor:
         """