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

CodePudding user response：

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)