tf.Module doesn't recognise non trainable variables from keras layers [TF 2.0]

[n3011]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 16.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device:
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): tf-nightly-2.0-preview==2.0.0.dev20190602
• Python version: 3.6
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version: CPU
• GPU model and memory:

When using Keras layers inside tf.Module module and setting trainable=False in the keras layer doesn't results in non-trainable variables in the tf.Module scope.
The below example code module M's trainable_variables should return 6, But it returns 8.

Code to reproduce the issue

class M(tf.Module):

    def __init__(self):
        super(M, self).__init__()
        self.list = []
        self.list.append([tf.keras.layers.Dense(5, trainable=False), tf.keras.layers.Dense(5)])
        self.list.append([tf.keras.layers.Dense(5), tf.keras.layers.Dense(5)])

    def __call__(self, inputs):
        output = inputs
        for l_list in self.list:
            for l in l_list:
                output = l(output)
        return output

m = M()
m(tf.ones((10, 10)))
Got: print(len(m.trainable_variables))  = 8

Expected: print(len(m.trainable_variables)) = 6

[gadagashwini-zz]

I could reproduce the reported issue here on tf-nightly-2.0-preview version. Thanks!

[qlzh727]

Thanks for reporting the issue. Here is some context about the root cause.

There are two "trainable" concept here, one is that whether the variable itself is trainable, the second is that whether the layer contains the variable is trainable. The first one is immutable, while the second is mutable. The layer/container's trainable state will affect the return value for trainable/non_trainable_variable.

When creating a variable, user could do layer.add_variable(trainable=False), that will force the variable to be non-trainable, regardless whether the layer itself is trainable or not.

In the case that the layer is not trainable, we will create the variable as trainable variable. We check layer.trainable state in the layer.trainable_weights/non_trainable_weights to return the correct value.

In your specific case, tf.module just recursively visit all its attribute, and find all the variable like object. It discards the container/layer trainable state, and relying on only the variable trainable information, which is why it returns 8 here.

A simple workaround is to change the base class for M to be layer, which will correctly check the sub layer trainable state, while we are fixing the issue on the tf.module side.

[qlzh727]

Also adding @tomhennigan who is the owner of tf.Module