tf.Module Explained: The Base for Layers and Models in TensorFlow

The Story — A Box That Knows Its Own Parts

In this module you will build a complete neural network from scratch: layers, activations, losses, optimizers and a training loop. No Keras layers, no model.fit(). Everything starts with tf.Module. (Keras layers and models are themselves built on top of it.)

Your First Module

A module is a Python class. Create variables in __init__. Do the maths in __call__.

import tensorflow as tf

class Scale(tf.Module):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.a = tf.Variable(2.0, name="a")                   # trainable
        self.b = tf.Variable(1.0, name="b")                   # trainable
        self.calls = tf.Variable(0, trainable=False, name="calls")

    def __call__(self, x):
        self.calls.assign_add(1)
        return self.a * x + self.b

s = Scale(name="scale")
print(s(tf.constant([1.0, 2.0, 3.0])))
print("all variables      :", [v.name for v in s.variables])
print("trainable variables:", [v.name for v in s.trainable_variables])

You never wrote a list of variables. The module found them by looking at its own attributes. The non-trainable counter is in variables but not in trainable_variables.

Modules Inside Modules

A model is a module that holds layer modules. The tracking goes all the way down.

class Linear(tf.Module):
    def __init__(self, in_dim, out_dim, name=None, trainable=True):
        super().__init__(name=name)
        with self.name_scope:                                  # prefix names with the module name
            self.w = tf.Variable(tf.random.normal([in_dim, out_dim]), name="w", trainable=trainable)
            self.b = tf.Variable(tf.zeros([out_dim]), name="b", trainable=trainable)

    def __call__(self, x):
        return x @ self.w + self.b

class TwoLayerNet(tf.Module):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.hidden = Linear(4, 8, name="hidden")
        self.out = Linear(8, 3, name="out")

    def __call__(self, x):
        return self.out(tf.nn.relu(self.hidden(x)))

net = TwoLayerNet(name="net")
print("output shape:", net(tf.ones([2, 4])).shape)
print("submodules  :", [m.name for m in net.submodules])
for v in net.trainable_variables:
    print(f"  {v.name:12s} {v.shape}")
print("parameters  :", sum(int(tf.size(v)) for v in net.trainable_variables))

Try It — Model Tree Builder

Fast Calls and Saving

Add tf.function

Decorate __call__ (or wrap the whole model) to run it as a graph. Everything from the last module still applies.

class FastNet(TwoLayerNet):
    @tf.function(input_signature=[tf.TensorSpec([None, 4], tf.float32)])
    def __call__(self, x):
        return super().__call__(x)

fast = FastNet(name="fast")
print(fast(tf.ones([5, 4])).shape)

Save and Restore Weights With a Checkpoint

ckpt = tf.train.Checkpoint(model=net)
path = ckpt.write("/tmp/net_ckpt")

before = net.out.b.numpy().copy()
net.out.b.assign([9.0, 9.0, 9.0])               # damage the weights
ckpt.read(path)                                   # restore them
print("restored:", (net.out.b.numpy() == before).all())

tf.Module vs Keras Layer

Golden Rules