Visualization approach 1

networks/modules.py

@@ -14,6 +14,9 @@ from torch.utils.data import DataLoader
 from dataset import DataContainer
 
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
 class LightningModuleOverrides:
 
     @property
@@ -25,8 +28,8 @@ class LightningModuleOverrides:
 
     @data_loader
     def tng_dataloader(self):
-        num_workers = 0  #  os.cpu_count() // 2
-        return DataLoader(DataContainer('data', self.size, self.step),
+        num_workers = 0  # os.cpu_count() // 2
+        return DataLoader(DataContainer(os.path.join('data', 'training'), self.size, self.step),
                           shuffle=True, batch_size=10000, num_workers=num_workers)
 
 
@@ -236,6 +239,19 @@ class Encoder(Module):
         return tensor
 
 
+class AttentionEncoder(Module):
+
+    def __init__(self):
+        super(AttentionEncoder, self).__init__()
+        self.l_stack = TimeDistributed(EncoderLinearStack())
+
+    def forward(self, x):
+        tensor = self.l_stack(x)
+        torch.bmm()  # TODO Add Attention here
+
+        return tensor
+
+
 class PoolingEncoder(Module):
 
     def __init__(self, lat_dim, variational=False):