During training the TensorFlow model(!!Not the Keras model), How to get the input and output of the intermediate layer(op) of the model?

Question

During training the TensorFlow model(!!Not the Keras model), is possible to get the input and output of the intermediate layer(op) of the model?

I use the example from Tensorflow tutorial as the demo:

To explain more clearly, I made a few changes to the model without changing its purpose.

Tensorflow Version: 2.8.0

class MyModel(Model):
  def __init__(self):
    super(MyModel, self).__init__()
    self.conv1 = Conv2D(32, 3, activation='relu')
    self.flatten = Flatten()
    self.d1 = Dense(128, activation='relu')
    self.d2 = Dense(10)

  def call(self, x):
    # x = self.conv1(x)
    # x = self.flatten(x)
    # x = self.d1(x)
    # return self.d2(x)
    x1 = self.conv1(x) 
    x2 = self.flatten(x1)
    x3 = self.d1(x2)
    return self.d2(x3)

Is possible to get the x1, x2, and x3 in the model or the input and output of self.conv1?

Answer 1

I would recommend using a custom Keras callback to feed data to the model during training and then saving the weights and outputs. You can feed the callback your training data or other data, for example your test data:

import tensorflow as tf
import numpy as np

class MyModel(tf.keras.Model):
  def __init__(self):
    super(MyModel, self).__init__()
    self.conv1 = tf.keras.layers.Conv2D(32, 3, activation='relu')
    self.flatten = tf.keras.layers.Flatten()
    self.d1 = tf.keras.layers.Dense(128, activation='relu')
    self.d2 = tf.keras.layers.Dense(10)

  def call(self, x):
    x1 = self.conv1(x) 
    x2 = self.flatten(x1)
    x3 = self.d1(x2)
    return self.d2(x3)

class CustomCallback(tf.keras.callbacks.Callback):
   def __init__(self, data):
        self.data = data
   def on_epoch_end(self, epoch, logs=None):
        #if epoch == some_value: <-- can also define a condition
        conv_layer = self.model.layers[0]
        outputs = conv_layer(self.data)
        np.save('conv_outputs', np.array(outputs)) 
        np.save('conv_weights', np.array(conv_layer.weights))
        tf.print('Saved Conv2D outputs and weights')

model = MyModel()
x_train = tf.random.normal((10, 32, 32, 3))
x_test = tf.random.normal((10, 32, 32, 3))
model.compile(optimizer='adam', loss = tf.keras.losses.SparseCategoricalCrossentropy(True))
model.fit(x_train, tf.random.uniform((10, 1), maxval=10), epochs=2, callbacks=[CustomCallback(x_test)], batch_size=2)