[architectures] modified the original DR-GAN to be as close as possible from...

[architectures] modified the original DR-GAN to be as close as possible from MSU (projection at the beginning of decoder)

bob/learn/pytorch/architectures/DRGANOriginal.py

@@ -157,6 +157,11 @@ class DRGANOriginal_decoder(nn.Module):
 
     self.ngpu = 1 # usually, we don't have more than one GPU 
 
+    self.project = nn.Sequential(
+      nn.Linear((self.noise_dim + self.conditional_dim + self.latent_dim), (self.latent_dim * 6 * 6)),
+      nn.Dropout(p=0.4),
+    )
+
     self.main = nn.Sequential(
       # input is 320x6x6, output is 160x6x6
       nn.ConvTranspose2d(320, 160, 3, 1, 1, bias=False),
@@ -237,41 +242,25 @@ class DRGANOriginal_decoder(nn.Module):
       The encoded ID for the minibatch 
     """    
     decoder_input = torch.cat((z, y, f), 1)
-    
-    
+    decoder_input = decoder_input.squeeze()
 
-    # linear transform to build a hypercube as input to deconv layers
-    #
-    # input is noise_dim + conditional_dim + latent_dim
-    # output is (latent_dim x 6 x 6)
-    # 
-    # Dropout + BatchNorm + ELU are applied to the cube
-    
-    # squeeze, apply linear layer, and unsqueeze
-    
-    # specify the ops
-    lin = nn.Linear((self.noise_dim + self.conditional_dim + self.latent_dim), (self.latent_dim * 6 * 6))
-    dropout = nn.Dropout(p=0.4)
+    # used in the "projection layer"
     bn = nn.BatchNorm2d(320)
     elu = nn.ELU(inplace=True)
     if torch.cuda.is_available():
-      lin = lin.cuda()
-      dropout = dropout.cuda()
       bn = bn.cuda()
       elu = elu.cuda()
 
-    decoder_input = torch.squeeze(decoder_input)
-    projected = lin(decoder_input)
-    projected = projected.unsqueeze(2)
-    projected = projected.unsqueeze(3)
-    dropped = dropout(projected)
-    reshaped = dropped.view(-1, self.latent_dim, 6, 6) 
-    hypercube = elu(bn(reshaped))
-
-    # deconv layers
+    # (projection + BN + ELU) + deconvolution 
     if isinstance(decoder_input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
+      projected = nn.parallel.data_parallel(self.project, decoder_input, range(self.ngpu))
+      hypercube = projected.view(-1, self.latent_dim, 6, 6)
+      hypercube = elu(bn(hypercube))
       output = nn.parallel.data_parallel(self.main, hypercube, range(self.ngpu))
     else:
+      projected = self.project(decoder_input)
+      hypercube = projected.view(-1, self.latent_dim, 6, 6)
+      hypercube = elu(bn(hypercube))
       output = self.main(hypercube)
    
     return output