Add Model02 recurrent generator: Model02RG

Author: Atcold

main.py

@@ -108,6 +108,8 @@ def main():
         from model.Model01 import Model01 as Model
     elif args.model == 'model_02':
         from model.Model02 import Model02 as Model
+    elif args.model == 'model_02_rg':
+        from model.Model02 import Model02RG as Model
     else:
         print('\n{:#^80}\n'.format(' Please select a valid model '))
         exit()
@@ -164,8 +166,14 @@ def adjust_learning_rate(opt, epoch):
 def selective_zero(s, new):
     if new.any():  # if at least one video changed
         b = new.nonzero().squeeze(1)  # get the list of indices
-        for layer in range(len(s)):  # for every layer having a state
-            s[layer] = s[layer].index_fill(0, V(b), 0)  # mask state, zero selected indices
+        if isinstance(s[0], list):  # recurrent G
+            for layer in range(len(s[0])):  # for every layer having a state
+                s[0][layer] = s[0][layer].index_fill(0, V(b), 0)  # mask state, zero selected indices
+            for layer in range(len(s[1])):  # for every layer having a state
+                s[1][layer] = s[1][layer].index_fill(0, V(b), 0)  # mask state, zero selected indices
+        else:  # simple convolutive G
+            for layer in range(len(s)):  # for every layer having a state
+                s[layer] = s[layer].index_fill(0, V(b), 0)  # mask state, zero selected indices
 
 
 def selective_match(x_hat, x, new):

model/Model02.py

@@ -6,6 +6,8 @@
 
 
 # Define some constants
+from model.RG import RG
+
 KERNEL_SIZE = 3
 PADDING = KERNEL_SIZE // 2
 KERNEL_STRIDE = 2
@@ -14,7 +16,7 @@
 
 class Model02(nn.Module):
     """
-    Generate a constructor for model_01 type of network
+    Generate a constructor for model_02 type of network
     """
 
     def __init__(self, network_size: tuple, input_spatial_size: tuple) -> None:
@@ -29,7 +31,7 @@ def __init__(self, network_size: tuple, input_spatial_size: tuple) -> None:
         super().__init__()
         self.hidden_layers = len(network_size) - 2
 
-        print('\n{:-^80}'.format(' Building model '))
+        print('\n{:-^80}'.format(' Building model Model02 '))
         print('Hidden layers:', self.hidden_layers)
         print('Net sizing:', network_size)
         print('Input spatial size: {} x {}'.format(network_size[0], input_spatial_size))
@@ -91,21 +93,114 @@ def forward(self, x, state):
         return (x, state), (x_mean, video_index)
 
 
-def _test_model():
+class Model02RG(nn.Module):
+    """
+    Generate a constructor for model_02_rg type of network
+    """
+
+    def __init__(self, network_size: tuple, input_spatial_size: tuple) -> None:
+        """
+        Initialise Model02RG constructor
+
+        :param network_size: (n, h1, h2, ..., emb_size, nb_videos)
+        :type network_size: tuple
+        :param input_spatial_size: (height, width)
+        :type input_spatial_size: tuple
+        """
+        super().__init__()
+        self.hidden_layers = len(network_size) - 2
+
+        print('\n{:-^80}'.format(' Building model Model02RG '))
+        print('Hidden layers:', self.hidden_layers)
+        print('Net sizing:', network_size)
+        print('Input spatial size: {} x {}'.format(network_size[0], input_spatial_size))
+
+        # main auto-encoder blocks
+        self.activation_size = [input_spatial_size]
+        for layer in range(0, self.hidden_layers):
+            # print some annotation when building model
+            print('{:-<80}'.format('Layer ' + str(layer + 1) + ' '))
+            print('Bottom size: {} x {}'.format(network_size[layer], self.activation_size[-1]))
+            self.activation_size.append(tuple(ceil(s / 2) for s in self.activation_size[layer]))
+            print('Top size: {} x {}'.format(network_size[layer + 1], self.activation_size[-1]))
+
+            # init D (discriminative) blocks
+            multiplier = layer and 2 or 1  # D_n, n > 1, has intra-layer feedback
+            setattr(self, 'D_' + str(layer + 1), nn.Conv2d(
+                in_channels=network_size[layer] * multiplier, out_channels=network_size[layer + 1],
+                kernel_size=KERNEL_SIZE, stride=KERNEL_STRIDE, padding=PADDING
+            ))
+            setattr(self, 'BN_D_' + str(layer + 1), nn.BatchNorm2d(network_size[layer + 1]))
+
+            # init G (generative) blocks
+            setattr(self, 'G_' + str(layer + 1), RG(
+                in_channels=network_size[layer + 1], out_channels=network_size[layer],
+                kernel_size=KERNEL_SIZE, stride=KERNEL_STRIDE, padding=PADDING
+            ))
+            setattr(self, 'BN_G_' + str(layer + 1), nn.BatchNorm2d(network_size[layer]))
+
+        # init auxiliary classifier
+        print('{:-<80}'.format('Classifier '))
+        print(network_size[-2], '-->', network_size[-1])
+        self.average = nn.AvgPool2d(self.activation_size[-1])
+        self.stabiliser = nn.Linear(network_size[-2], network_size[-1])
+        print(80 * '-', end='\n\n')
+
+    def forward(self, x, state):
+        activation_sizes = [x.size()]  # start from the input
+        residuals = list()
+        # state[0] --> network layer state; state[1] --> generative state
+        state = state or [[None] * (self.hidden_layers - 1), [None] * self.hidden_layers]
+        for layer in range(0, self.hidden_layers):  # connect discriminative blocks
+            if layer:  # concat the input with the state for D_n, n > 1
+                s = state[0][layer - 1] or V(x.data.clone().zero_())
+                x = torch.cat((x, s), 1)
+            x = getattr(self, 'D_' + str(layer + 1))(x)
+            residuals.append(x)
+            x = f.relu(x)
+            x = getattr(self, 'BN_D_' + str(layer + 1))(x)
+            activation_sizes.append(x.size())  # cache output size for later retrieval
+        for layer in reversed(range(0, self.hidden_layers)):  # connect generative blocks
+            x = getattr(self, 'G_' + str(layer + 1))((x, activation_sizes[layer]), state[1][layer])
+            state[1][layer] = x  # h[t - 1] <- h[t]
+            if layer:
+                state[0][layer - 1] = x
+                x += residuals[layer - 1]
+            x = f.relu(x)
+            x = getattr(self, 'BN_G_' + str(layer + 1))(x)
+        x_mean = self.average(residuals[-1])
+        video_index = self.stabiliser(x_mean.view(x_mean.size(0), -1))
+
+        return (x, state), (x_mean, video_index)
+
+
+def _test_models():
+    _test_model(Model02)
+    _test_model(Model02RG)
+
+
+def _test_model(Model):
     big_t = 2
     x = torch.rand(big_t + 1, 1, 3, 4 * 2**3 + 3, 6 * 2**3 + 5)
     big_k = 10
     y = torch.LongTensor(big_t, 1).random_(big_k)
-    model_01 = Model02(network_size=(3, 6, 12, 18, big_k), input_spatial_size=x[0].size()[2:])
+    model = Model(network_size=(3, 6, 12, 18, big_k), input_spatial_size=x[0].size()[2:])
 
     state = None
-    (x_hat, state), (emb, idx) = model_01(V(x[0]), state)
+    (x_hat, state), (emb, idx) = model(V(x[0]), state)
 
     print('Input size:', tuple(x.size()))
     print('Output size:', tuple(x_hat.data.size()))
     print('Video index size:', tuple(idx.size()))
     for i, s in enumerate(state):
-        print('State', i + 1, 'has size:', tuple(s.size()))
+        if isinstance(s, list):
+            for i, s in enumerate(state[0]):
+                print('Net state', i + 1, 'has size:', tuple(s.size()))
+            for i, s in enumerate(state[1]):
+                print('G', i + 1, 'state has size:', tuple(s.size()))
+            break
+        else:
+            print('State', i + 1, 'has size:', tuple(s.size()))
     print('Embedding has size:', emb.data.numel())
 
     mse = nn.MSELoss()
@@ -118,7 +213,7 @@ def _test_model():
     show_graph(loss_t1)
 
     # run one more time
-    (x_hat, _), (_, idx) = model_01(V(x[1]), state)
+    (x_hat, _), (_, idx) = model(V(x[1]), state)
 
     x_next = V(x[2])
     y_var = V(y[1])
@@ -128,7 +223,12 @@ def _test_model():
     show_graph(loss_tot)
 
 
-def _test_training():
+def _test_training_models():
+    _test_training(Model02)
+    _test_training(Model02RG)
+
+
+def _test_training(Model):
     big_k = 10  # number of training videos
     network_size = (3, 6, 12, 18, big_k)
     big_t = 6  # sequence length
@@ -147,7 +247,7 @@ def _test_training():
     print('Target index has size', tuple(y.size()))
 
     print('Define model')
-    model = Model02(network_size=network_size, input_spatial_size=x[0].size()[2:])
+    model = Model(network_size=network_size, input_spatial_size=x[0].size()[2:])
 
     print('Create a MSE and NLL criterions')
     mse = nn.MSELoss()
@@ -175,13 +275,13 @@ def _test_training():
 
 
 if __name__ == '__main__':
-    _test_model()
-    _test_training()
+    _test_models()
+    _test_training_models()
 
 
 __author__ = "Alfredo Canziani"
 __credits__ = ["Alfredo Canziani"]
 __maintainer__ = "Alfredo Canziani"
 __email__ = "alfredo.canziani@gmail.com"
 __status__ = "Production"  # "Prototype", "Development", or "Production"
-__date__ = "Feb 17"
+__date__ = "Feb, Mar 17"

model/RG.py

@@ -0,0 +1,32 @@
+from torch import nn
+
+
+class RG(nn.Module):
+    """Recurrent Generative Module"""
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
+        """ Initialise RG Module (parameters as nn.ConvTranspose2d)"""
+        super().__init__()
+        self.from_input = nn.ConvTranspose2d(
+            in_channels=in_channels, out_channels=out_channels,
+            kernel_size=kernel_size, stride=stride, padding=padding
+        )
+        self.from_state = nn.Conv2d(
+            in_channels=out_channels, out_channels=out_channels,
+            kernel_size=kernel_size, padding=padding, bias=False
+        )
+
+    def forward(self, x, state):
+        """
+        Calling signature
+
+        :param x: (input, output_size)
+        :type x: tuple
+        :param state: previous output
+        :type state: torch.Tensor
+        :return: current state
+        :rtype: torch.Tensor
+        """
+        x = self.from_input(*x)  # the very first x is a tuple (input, expected_output_size)
+        if state: x += self.from_state(state)
+        return x