VAE Debugged and Running

lib/modules/blocks.py

@@ -17,9 +17,9 @@ class ConvModule(nn.Module):
         output = self(x)
         return output.shape[1:]
 
-    def __init__(self, in_shape, activation: nn.Module = nn.ELU, pooling_size=None, use_bias=True, use_norm=False,
-                 dropout: Union[int, float] = 0, conv_class=nn.Conv2d,
-                 conv_filters=64, conv_kernel=5, conv_stride=1, conv_padding=0):
+    def __init__(self, in_shape, conv_filters, conv_kernel, activation: nn.Module = nn.ELU, pooling_size=None,
+                 use_bias=True, use_norm=False, dropout: Union[int, float] = 0,
+                 conv_class=nn.Conv2d, conv_stride=1, conv_padding=0):
         super(ConvModule, self).__init__()
 
         # Module Parameters
@@ -30,12 +30,14 @@ class ConvModule(nn.Module):
         # Convolution Parameters
         self.padding = conv_padding
         self.stride = conv_stride
+        self.conv_filters = conv_filters
+        self.conv_kernel = conv_kernel
 
         # Modules
         self.dropout = nn.Dropout2d(dropout) if dropout else lambda x: x
         self.pooling = nn.MaxPool2d(pooling_size) if pooling_size else lambda x: x
-        self.norm = nn.BatchNorm2d(in_channels, eps=1e-04, affine=False) if use_norm else lambda x: x
-        self.conv = conv_class(in_channels, conv_filters, conv_kernel, bias=use_bias,
+        self.norm = nn.BatchNorm2d(in_channels, eps=1e-04) if use_norm else lambda x: x
+        self.conv = conv_class(in_channels, self.conv_filters, self.conv_kernel, bias=use_bias,
                                padding=self.padding, stride=self.stride
                                )
 
@@ -57,22 +59,23 @@ class DeConvModule(nn.Module):
         output = self(x)
         return output.shape[1:]
 
-    def __init__(self, in_shape, conv_filters=3, conv_kernel=5, conv_stride=1, conv_padding=0,
-                 dropout: Union[int, float] = 0, autopad=False,
-                 activation: Union[None, nn.Module] = nn.ReLU, interpolation_scale=None,
+    def __init__(self, in_shape, conv_filters, conv_kernel, conv_stride=1, conv_padding=0,
+                 dropout: Union[int, float] = 0, autopad=0,
+                 activation: Union[None, nn.Module] = nn.ReLU, interpolation_scale=0,
                  use_bias=True, use_norm=False):
         super(DeConvModule, self).__init__()
         in_channels, height, width = in_shape[0], in_shape[1], in_shape[2]
         self.padding = conv_padding
+        self.conv_kernel = conv_kernel
         self.stride = conv_stride
         self.in_shape = in_shape
         self.conv_filters = conv_filters
 
         self.autopad = AutoPad() if autopad else lambda x: x
         self.interpolation = Interpolate(scale_factor=interpolation_scale) if interpolation_scale else lambda x: x
-        self.norm = nn.BatchNorm2d(in_channels, eps=1e-04, affine=False) if use_norm else lambda x: x
+        self.norm = nn.BatchNorm2d(in_channels, eps=1e-04) if use_norm else lambda x: x
         self.dropout = nn.Dropout2d(dropout) if dropout else lambda x: x
-        self.de_conv = nn.ConvTranspose2d(in_channels, self.conv_filters, conv_kernel, bias=use_bias,
+        self.de_conv = nn.ConvTranspose2d(in_channels, self.conv_filters, self.conv_kernel, bias=use_bias,
                                           padding=self.padding, stride=self.stride)
 
         self.activation = activation() if activation else lambda x: x

lib/modules/utils.py

@@ -6,8 +6,6 @@ from torch import nn
 from torch import functional as F
 from torch.utils.data import DataLoader
 
-from lib.objects.map import MapStorage
-
 import pytorch_lightning as pl
 
 
@@ -27,10 +25,11 @@ class Flatten(nn.Module):
             print(e)
             return -1
 
-    def __init__(self, in_shape, to=(-1, )):
+    def __init__(self, in_shape, to=-1):
+        assert isinstance(to, int) or isinstance(to, tuple)
         super(Flatten, self).__init__()
         self.in_shape = in_shape
-        self.to = to
+        self.to = (to,) if isinstance(to, int) else to
 
     def forward(self, x):
         return x.view(x.size(0), *self.to)
@@ -107,7 +106,7 @@ class LightningBaseModule(pl.LightningModule, ABC):
         # Data loading
         # =============================================================================
         # Map Object
-        self.map_storage = MapStorage(self.hparams.data_param.map_root)
+        # self.map_storage = MapStorage(self.hparams.data_param.map_root)
 
     def size(self):
         return self.shape