Add alternative implementation of tile_concat that uses much less weights and memory.

Author: alexlee-gk

video_prediction/models/savp_model.py

@@ -357,7 +357,10 @@ def _rnn_func(self, inputs, state, num_units):
         return rnn_cell(inputs, state)
 
     def _conv_rnn_func(self, inputs, state, filters):
-        inputs_shape = inputs.get_shape().as_list()
+        if isinstance(inputs, (list, tuple)):
+            inputs_shape = inputs[0].shape.as_list()
+        else:
+            inputs_shape = inputs.shape.as_list()
         input_shape = inputs_shape[1:]
         if self.hparams.conv_rnn_norm_layer == 'none':
             normalizer_fn = None
@@ -446,18 +449,26 @@ def concat(tensors, axis):
                     h = layers[-1][-1]
                     kernel_size = (3, 3)
                 if self.hparams.where_add == 'all' or (self.hparams.where_add == 'input' and i == 0):
-                    h = tile_concat([h, state_action_z[:, None, None, :]], axis=-1)
-                h = downsample_layer(h, out_channels, kernel_size=kernel_size, strides=(2, 2))
+                    if self.hparams.use_tile_concat:
+                        h = tile_concat([h, state_action_z[:, None, None, :]], axis=-1)
+                    else:
+                        h = [h, state_action_z]
+                h = _maybe_tile_concat_layer(downsample_layer)(
+                    h, out_channels, kernel_size=kernel_size, strides=(2, 2))
                 h = norm_layer(h)
                 h = activation_layer(h)
             if use_conv_rnn:
                 with tf.variable_scope('%s_h%d' % ('conv' if self.hparams.ablation_rnn else self.hparams.conv_rnn, i)):
                     if self.hparams.where_add == 'all':
-                        conv_rnn_h = tile_concat([h, state_action_z[:, None, None, :]], axis=-1)
+                        if self.hparams.use_tile_concat:
+                            conv_rnn_h = tile_concat([h, state_action_z[:, None, None, :]], axis=-1)
+                        else:
+                            conv_rnn_h = [h, state_action_z]
                     else:
                         conv_rnn_h = h
                     if self.hparams.ablation_rnn:
-                        conv_rnn_h = conv2d(conv_rnn_h, out_channels, kernel_size=(5, 5))
+                        conv_rnn_h = _maybe_tile_concat_layer(conv2d)(
+                            conv_rnn_h, out_channels, kernel_size=(5, 5))
                         conv_rnn_h = norm_layer(conv_rnn_h)
                         conv_rnn_h = activation_layer(conv_rnn_h)
                     else:
@@ -474,18 +485,25 @@ def concat(tensors, axis):
                 else:
                     h = tf.concat([layers[-1][-1], layers[num_encoder_layers - i - 1][-1]], axis=-1)
                 if self.hparams.where_add == 'all' or (self.hparams.where_add == 'middle' and i == 0):
-                    h = tile_concat([h, state_action_z[:, None, None, :]], axis=-1)
-                h = upsample_layer(h, out_channels, kernel_size=(3, 3), strides=(2, 2))
+                    if self.hparams.use_tile_concat:
+                        h = tile_concat([h, state_action_z[:, None, None, :]], axis=-1)
+                    else:
+                        h = [h, state_action_z]
+                h = _maybe_tile_concat_layer(upsample_layer)(
+                    h, out_channels, kernel_size=(3, 3), strides=(2, 2))
                 h = norm_layer(h)
                 h = activation_layer(h)
             if use_conv_rnn:
                 with tf.variable_scope('%s_h%d' % ('conv' if self.hparams.ablation_rnn else self.hparams.conv_rnn, len(layers))):
                     if self.hparams.where_add == 'all':
-                        conv_rnn_h = tile_concat([h, state_action_z[:, None, None, :]], axis=-1)
+                        if self.hparams.use_tile_concat:
+                            conv_rnn_h = tile_concat([h, state_action_z[:, None, None, :]], axis=-1)
+                        else:
+                            conv_rnn_h = [h, state_action_z]
                     else:
                         conv_rnn_h = h
                     if self.hparams.ablation_rnn:
-                        conv_rnn_h = conv2d(conv_rnn_h, out_channels, kernel_size=(5, 5))
+                        conv_rnn_h = _maybe_tile_concat_layer(conv2d)(conv_rnn_h, out_channels, kernel_size=(5, 5))
                         conv_rnn_h = norm_layer(conv_rnn_h)
                         conv_rnn_h = activation_layer(conv_rnn_h)
                     else:
@@ -770,6 +788,7 @@ def get_default_hparams_dict(self):
             kernel_size=(5, 5),
             dilation_rate=(1, 1),
             where_add='all',
+            use_tile_concat=True,
             learn_initial_state=False,
             rnn='lstm',
             conv_rnn='lstm',
@@ -950,3 +969,16 @@ def center_slice(k):
     kernel[center_slice(kh), center_slice(kw)] = 1.0
     kernel /= np.sum(kernel)
     return kernel
+
+
+def _maybe_tile_concat_layer(conv2d_layer):
+    def layer(inputs, out_channels, *args, **kwargs):
+        if isinstance(inputs, (list, tuple)):
+            inputs_spatial, inputs_non_spatial = inputs
+            outputs = (conv2d_layer(inputs_spatial, out_channels, *args, **kwargs) +
+                       dense(inputs_non_spatial, out_channels, use_bias=False)[:, None, None, :])
+        else:
+            outputs = conv2d_layer(inputs, out_channels, *args, **kwargs)
+        return outputs
+
+    return layer

video_prediction/ops.py

@@ -2,15 +2,17 @@
 import tensorflow as tf
 
 
-def dense(inputs, units, use_spectral_norm=False):
+def dense(inputs, units, use_spectral_norm=False, use_bias=True):
     with tf.variable_scope('dense'):
         input_shape = inputs.get_shape().as_list()
         kernel_shape = [input_shape[1], units]
         kernel = tf.get_variable('kernel', kernel_shape, dtype=tf.float32, initializer=tf.truncated_normal_initializer(stddev=0.02))
         if use_spectral_norm:
             kernel = spectral_normed_weight(kernel)
-        bias = tf.get_variable('bias', [units], dtype=tf.float32, initializer=tf.zeros_initializer())
-        outputs = tf.matmul(inputs, kernel) + bias
+        outputs = tf.matmul(inputs, kernel)
+        if use_bias:
+            bias = tf.get_variable('bias', [units], dtype=tf.float32, initializer=tf.zeros_initializer())
+            outputs = tf.nn.bias_add(outputs, bias)
         return outputs
 
 

video_prediction/rnn_ops.py

@@ -125,11 +125,25 @@ def _conv2d(self, inputs):
             outputs = nn_ops.bias_add(outputs, bias)
         return outputs
 
+    def _dense(self, inputs):
+        num_units = 4 * self._filters
+        input_shape = inputs.shape.as_list()
+        kernel_shape = [input_shape[-1], num_units]
+        kernel = vs.get_variable("weights", kernel_shape, dtype=dtypes.float32,
+                                 initializer=init_ops.truncated_normal_initializer(stddev=0.02))
+        outputs = tf.matmul(inputs, kernel)
+        return outputs
+
     def call(self, inputs, state):
         """2D Convolutional LSTM cell with (optional) normalization and recurrent dropout."""
         c, h = state
+        tile_concat = isinstance(inputs, (list, tuple))
+        if tile_concat:
+            inputs, inputs_non_spatial = inputs
         args = array_ops.concat([inputs, h], -1)
         concat = self._conv2d(args)
+        if tile_concat:
+            concat = concat + self._dense(inputs_non_spatial)[:, None, None, :]
 
         if self._normalizer_fn and not self._separate_norms:
             concat = self._norm(concat, "input_transform_forget_output")
@@ -209,13 +223,26 @@ def _conv2d(self, inputs, output_filters, bias_initializer):
             outputs = nn_ops.bias_add(outputs, bias)
         return outputs
 
+    def _dense(self, inputs, num_units):
+        input_shape = inputs.shape.as_list()
+        kernel_shape = [input_shape[-1], num_units]
+        kernel = vs.get_variable("weights", kernel_shape, dtype=dtypes.float32,
+                                 initializer=init_ops.truncated_normal_initializer(stddev=0.02))
+        outputs = tf.matmul(inputs, kernel)
+        return outputs
+
     def call(self, inputs, state):
         bias_ones = self._bias_initializer
         if self._bias_initializer is None:
             bias_ones = init_ops.ones_initializer()
+        tile_concat = isinstance(inputs, (list, tuple))
+        if tile_concat:
+            inputs, inputs_non_spatial = inputs
         with vs.variable_scope('gates'):
             inputs = array_ops.concat([inputs, state], axis=-1)
             concat = self._conv2d(inputs, 2 * self._filters, bias_ones)
+            if tile_concat:
+                concat = concat + self._dense(inputs_non_spatial, concat.shape[-1].value)[:, None, None, :]
             if self._normalizer_fn and not self._separate_norms:
                 concat = self._norm(concat, "reset_update", bias_ones)
             r, u = array_ops.split(concat, 2, axis=-1)
@@ -230,6 +257,8 @@ def call(self, inputs, state):
         with vs.variable_scope('candidate'):
             inputs = array_ops.concat([inputs, r * state], axis=-1)
             candidate = self._conv2d(inputs, self._filters, bias_zeros)
+            if tile_concat:
+                candidate = candidate + self._dense(inputs_non_spatial, candidate.shape[-1].value)[:, None, None, :]
             if self._normalizer_fn:
                 candidate = self._norm(candidate, "state", bias_zeros)