tensorflow2-recompute-grad-testing.py

import tensorflow as tf


tf.random.set_seed(20)


class Layer(tf.keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.dense = tf.keras.layers.Dense(20)
        self.dropout1 = tf.keras.layers.Dropout(0.1, seed=20)
        self.dense2 = tf.keras.layers.Dense(20)

    def call(self, input_tensor):
        return self.dense2(self.dropout1(self.dense(input_tensor)))


class Model(tf.keras.Model):
    def __init__(self):
        super().__init__()
        self.dense = tf.keras.layers.Dense(20)
        self.dropout1 = tf.keras.layers.Dropout(0.1, seed=20)
        self.dense2 = Layer()

    def call(self, input_tensor):
        return self.dense2(self.dropout1(self.dense(input_tensor)))


model = Model()


@tf.recompute_grad
def f(input_tensor):
    return model(input_tensor)


random_input = tf.random.uniform((20, 30))
with tf.GradientTape() as g:
    g.watch(model.trainable_variables)
    output = f(random_input)
checkpointed = g.gradient(output, model.trainable_variables)
print(checkpointed)

with tf.GradientTape() as g:
    g.watch(model.trainable_variables)
    output = model(random_input)
original = g.gradient(output, model.trainable_variables)
print(original)

for c, o in zip(checkpointed, original):
    rtol = tf.abs((c - o) / c)
    atol = tf.abs(c - o)

    print(rtol, atol)