Indicating training/testing modes in sonnet callbacks

What is a convenient way of providing boolean training flags, e.g. is_training that indicate, for example, whether to use batch_norm or not when using sonnet callback functions?

Example:

def make_transpose_cnn_model():
    def transpose_convnet1d(inputs):
        inputs = tf.expand_dims(inputs, axis=2)

        outputs = snt.Conv1DTranspose(output_channels=2, kernel_shape=10, stride=1)(inputs)
        outputs = snt.BatchNorm()(outputs, is_training=True) <- want to have this as input
        outputs = tf.nn.relu(outputs)
        outputs = snt.Conv1DTranspose(output_channels=2, kernel_shape=10, stride=1)(outputs)
        outputs = snt.BatchNorm()(outputs, is_training=True) <- want to have this as input
        outputs = tf.nn.relu(outputs)
        outputs = snt.BatchFlatten()(outputs)
        #outputs = tf.nn.dropout(outputs, keep_prob=tf.constant(1.0)) <- want to have this as input
        outputs = snt.Linear(output_size=128)(outputs)

        return outputs

    return transpose_convnet1d`

and

self._network = modules.GraphIndependent(
                edge_model_fn=EncodeProcessDecode.make_mlp_model,
                node_model_fn=EncodeProcessDecode.make_transpose_cnn_model,
                global_model_fn = EncodeProcessDecode.make_mlp_model)

I can’t pass this parameter in the _build() function as shown in the following since the interface of modules.GraphIndipendent won’t allow it:

    def _build(self, input_op, num_processing_steps, is_training=True):
        latent = self._encoder(input_op, is_training)

it yields:

TypeError: _build() got an unexpected keyword argument ‘is_training’

[malcolmreynolds]

If you have a module that you want to compose out of other modules, and some of the submodules require extra arguments like is_training, the canonical thing to do would be to define a new subclass of AbstractModule, rather than writing a function as you have. In particular, defining a module allows you to explicitly reuse variables by just calling the module twice (potentially with different is_training kwarg values).

Your example would be something like this:

class TransposeCnnModel(snt.AbstractModule):
  def __init__(self, name='transpose_cnn_model'):
    super(TransposeCnnModel, self).__init__(name=name)

  def _build(self, inputs, is_training):
    inputs = tf.expand_dims(inputs, axis=2)

    outputs = snt.Conv1DTranspose(output_channels=2, kernel_shape=10, stride=1)(inputs)
    outputs = snt.BatchNorm()(outputs, is_training=is_training)
    outputs = tf.nn.relu(outputs)
    outputs = snt.Conv1DTranspose(output_channels=2, kernel_shape=10, stride=1)(outputs)
    outputs = snt.BatchNorm()(outputs, is_training=is_training)
    outputs = tf.nn.relu(outputs)
    outputs = snt.BatchFlatten()(outputs)
    keep_prob = 0.7 if is_training else 1.0
    outputs = tf.nn.dropout(outputs, keep_prob=keep_prob)
    outputs = snt.Linear(output_size=128)(outputs)

    return outputs    

You would probably want to make more things configurable as constructor args (e.g. final output size, channels in the conv, dropout probability for training etc) but the above should fit the API that the GraphNet library expects.

As to your second point, I’m not sure exactly what you mean by Sonnet callback functions - could you be more specific?

[ferreirafabio]

it worked! Ty