How can you compute the average weights of sequentially saved checkpoints within a Subclassing Model using TensorFlow2?

edit

One way to compute the average weights of sequentially saved checkpoints within a Subclassing Model using TensorFlow2 is to load each checkpoint one by one and accumulate the weights. Once all the checkpoints are loaded, the accumulated weights can be divided by the number of checkpoints to get the average weights.

Here's an example code snippet to do that:

import tensorflow as tf

class SubclassingModel(tf.keras.Model):
    def __init__(self):
        super(SubclassingModel, self).__init__()
        self.dense1 = tf.keras.layers.Dense(64, activation='relu')
        self.dense2 = tf.keras.layers.Dense(10)

    def call(self, inputs):
        x = self.dense1(inputs)
        return self.dense2(x)

model = SubclassingModel()

# Define optimizer and loss function
optimizer = tf.keras.optimizers.Adam()
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# Train the model and save checkpoints
for epoch in range(EPOCHS):
    # Train step here

    # Save checkpoint
    model.save_weights('checkpoint_{}'.format(epoch))

# Load all checkpoints and compute average weights
num_checkpoints = 5
accumulated_weights = None
for i in range(num_checkpoints):
    model.load_weights('checkpoint_{}'.format(i))
    if i == 0:
        accumulated_weights = model.get_weights()
    else:
        accumulated_weights = [aw + w for aw, w in zip(accumulated_weights, model.get_weights())]

average_weights = [aw / num_checkpoints for aw in accumulated_weights]

# Use the average weights to initialize the model
model.set_weights(average_weights)