Merge pull request #41 from Tensor46/develop

Develop

Author: Tensor46

README.md

@@ -104,6 +104,7 @@ Trained on CIFAR10 (pggan-cifar10.py) -- requires more training (more gpus)!
   * [ReductionB](https://arxiv.org/pdf/1602.07261.pdf)
   * [ContextNet_Bottleneck](https://arxiv.org/pdf/1805.04554.pdf)
 
+* [DropBlock](https://arxiv.org/abs/1810.12890)
 * [PrimaryCapsule](https://arxiv.org/pdf/1710.09829.pdf)
 * [RoutingCapsule](https://arxiv.org/pdf/1710.09829.pdf)
 

core/NeuralArchitectures/convolutionalvae.py

@@ -10,43 +10,36 @@
 import numpy as np
 
 
-class ReShape(nn.Module):
-    def __init__(self, tensor_size):
-        super(ReShape, self).__init__()
-        self.tensor_size = tensor_size
-
-    def forward(self, tensor):
-        return tensor.view(tensor.size(0), *self.tensor_size[1:])
-
-
 class ConvolutionalVAE(nn.Module):
-    """
-        Convolutional Variational Auto Encoder
-
-        Parameters
-            tensor_size :: expected size of input tensor
-            embedding_layers :: a list of (filter_size, out_channels, strides)
-                                in each intermediate layer of the encoder.
-                                A flip is used for decoder
-            n_latent :: length of latent vecotr Z
-            decoder_final_activation :: tanh/sigm
-
-            activation, normalization, pre_nm, weight_nm, equalized, bias ::
-                refer to core.NeuralLayers
+    r""" Example Convolutional Variational Auto Encoder
+
+    Args:
+        tensor_size: shape of tensor in BCHW
+            (None/any integer >0, channels, height, width)
+        embedding_layers: a list of (filter_size, out_channels, strides)
+            in each intermediate layer of the encoder. A flip is used for
+            decoder.
+        n_latent: length of latent vecotr Z
+        decoder_final_activation: tanh/sigm
+        activation, normalization, pre_nm, weight_nm, equalized, bias:
+                refer to core.NeuralLayers.Convolution
+
+    Return:
+        encoded, mu, log_var, latent, decoded, kld, mse
     """
     def __init__(self,
-                 tensor_size=(6, 1, 28, 28),
-                 embedding_layers=[(3, 32, 2), (3, 64, 2)],
-                 n_latent=128,
-                 decoder_final_activation="tanh",
-                 pad=True,
-                 activation="relu",
-                 normalization=None,
-                 pre_nm=False,
-                 groups=1,
-                 weight_nm=False,
-                 equalized=False,
-                 bias=False,
+                 tensor_size: tuple = (6, 1, 28, 28),
+                 embedding_layers: list = [(3, 32, 2), (3, 64, 2)],
+                 n_latent: int = 128,
+                 decoder_final_activation: str = "tanh",
+                 pad: bool = True,
+                 activation: str = "relu",
+                 normalization: str = None,
+                 pre_nm: bool = False,
+                 groups: int = 1,
+                 weight_nm: bool = False,
+                 equalized: bool = False,
+                 bias: bool = False,
                  *args, **kwargs):
         super(ConvolutionalVAE, self).__init__()
 
@@ -65,23 +58,22 @@ def __init__(self,
         kwargs["equalized"] = equalized
         # encoder with Convolution layers
         encoder = []
-        _tensor_size = tensor_size
+        t_size = tensor_size
         for f, c, s in embedding_layers:
-            encoder.append(Convolution(_tensor_size, f, c, s, **kwargs))
-            _tensor_size = encoder[-1].tensor_size
+            encoder.append(Convolution(t_size, f, c, s, **kwargs))
+            t_size = encoder[-1].tensor_size
         self.encoder = nn.Sequential(*encoder)
 
         # mu and log_var to synthesize Z
-        self.mu = Linear(_tensor_size, n_latent, "", 0., bias=bias)
-        self.log_var = Linear(_tensor_size, n_latent, "", 0., bias=bias)
+        self.mu = Linear(t_size, n_latent, "", 0., bias=bias)
+        self.log_var = Linear(t_size, n_latent, "", 0., bias=bias)
 
         # decoder - (Linear layer + ReShape) to generate encoder last output
         # shape, followed by inverse of encoder
         decoder = []
         decoder.append(Linear(self.mu.tensor_size,
-                              int(np.prod(_tensor_size[1:])),
-                              activation, 0., bias=bias))
-        decoder.append(ReShape(_tensor_size))
+                              int(np.prod(t_size[1:])), activation, 0.,
+                              bias=bias, out_shape=t_size[1:]))
 
         decoder_layers = []
         for i, x in enumerate(embedding_layers[::-1]):
@@ -94,10 +86,9 @@ def __init__(self,
         for i, (f, c, s, o) in enumerate(decoder_layers):
             if i == len(decoder_layers)-1:
                 kwargs["activation"] = None
-            decoder.append(Convolution(_tensor_size, f, c, s,
-                                       transpose=True, **kwargs))
-            decoder[-1].tensor_size = o  # adjusting the output tensor size
-            _tensor_size = decoder[-1].tensor_size
+            decoder.append(Convolution(t_size, f, c, s, transpose=True,
+                                       maintain_out_size=True, **kwargs))
+            t_size = decoder[-1].tensor_size
         self.decoder = nn.Sequential(*decoder)
 
         # Final normalization
@@ -123,6 +114,7 @@ def forward(self, tensor, noisy_tensor=None):
         mse = F.mse_loss(decoded, tensor)
         return encoded, mu, log_var, latent, decoded, kld, mse
 
+
 # from core.NeuralLayers import Convolution, Linear
 # tensor_size = (1, 1, 28, 28)
 # tensor = torch.rand(*tensor_size)

core/NeuralEssentials/cudamodel.py

@@ -35,11 +35,13 @@ def check_precision_device(self, inputs):
             self.precision = p.dtype if "p" in locals() else torch.float32
         if type(inputs) in [list, tuple]:
             if self.is_cuda:
-                inputs = [x.type(self.precision).cuda() if self.is_cuda else
-                          x.type(self.precision) for x in inputs]
+                inputs = [(x.type(self.precision).cuda() if self.is_cuda else
+                           x.type(self.precision)) if x.dtype != torch.long
+                          else x for x in inputs]
             return inputs
         else:
-            inputs = inputs.type(self.precision)
+            if not (inputs.dtype == torch.long):
+                inputs = inputs.type(self.precision)
             if self.is_cuda:
                 inputs = inputs.cuda()
         return inputs