tf.keras.Model doesn't find assigned variables

Viewed 632

Look at this code (assume Tensorflow 2.1 or 2.2):

import tensorflow as tf

inputs = tf.keras.Input([])
var = tf.Variable(3.0)
out = var*inputs

model = tf.keras.Model(inputs=inputs, outputs=out)
print(model.variables)             # prints []

model.saved_vars = tf.Module()
model.saved_vars.var = var

print(model.saved_vars.variables)  # prints (<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=3.0>,)
print(model.variables)             # prints []

Why is model.variables empy? I would expect to contain var. As I understant the tf.Module, it works as a namescope and also records tf.Variable assigned to their attributes. Since a tf.keras.Model is a tf.Module, it should do the same recursively, shouldn't it?

1 Answers

As per the documentation, variables are Sequence of variables owned by this module and its submodules. But when we tested in Tensorflow version 2.2.0, the model.variables is not displaying variables of submodules.

Variables are automatically tracked when assigned to attributes of types inheriting from tf.Module. As you are doing model.saved_vars = tf.Module(), the model.saved_vars submodule becomes a named container for tf.Variables. Since tf.keras.Model is by default a tf.Module, you can assign var directly. But as you rightly mentioned the model.variables must record tf.Variables of submodule model.saved_vars.

Code - Tried adding another submodule inside model.saved_vars to check if print(model.saved_vars.variables) gives submodule variables. But still No.

import tensorflow as tf

inputs_1 = tf.keras.Input(shape=(3,))
var = tf.Variable(3.0)
out = var*inputs_1

model = tf.keras.Model(inputs=inputs_1, outputs=out)
print(model.variables)             # prints []

# Create Module
model.saved_vars = tf.Module()
model.saved_vars.var = var

model.var = var

print(model.saved_vars.variables)
print(model.variables)  

# Create Sub Module
model.saved_vars.saved_vars2 = tf.Module()
model.saved_vars.saved_vars2.var2 = var

print(model.saved_vars.variables)
print(model.variables)  

Output -

[]
(<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=3.0>,)
[<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=3.0>]
(<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=3.0>,)
[<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=3.0>]

Will share more information as and when we get.

Related