Getting NN weights for every batch / epoch from Keras model

Question

I am trying to get weights for every batch / epoch from Keras model after it is trained. To do so I use callback to make model save weights during training. Yet after model is trained it looks like I get weights only from the final epoch. How to get all weights that model generates? Here is a simple example:

import numpy as np
import tensorflow as tf
from tensorflow import keras
from keras import layers

# Generate data

start, stop = 1,100
cnt = stop - start + 1
xs = np.linspace(start, stop, num = cnt)
b,k = 1,2
ys = np.array([k*x + b for x in xs])

# Simple model with one feature and one unit for regression task

model = keras.Sequential([
    layers.Dense(units=1, input_shape=[1], activation='relu')
])
model.compile(loss='mae', optimizer='adam')
batch_size = int(cnt / 5)
epochs = 80

Next goes callback to save the Keras model weights at some frequency. According to Keras docs:

save_freq: 'epoch' or integer. When using 'epoch', the callback should save the model after each epoch. When using integer, the callback should save the model at end of this many batches.

checkpoint_filepath = './checkpoint.hdf5'
model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_filepath,
    save_weights_only=True,
    save_freq ='epoch', # 1 for every batch
    save_best_only=False
)

# Train model

history = model.fit(xs, ys, batch_size=batch_size, epochs=epochs, 
                    callbacks=[model_checkpoint_callback])

I use two different ways to get weights. First:

w, b = model.weights
print("Weights: \n {} \n Bias: \n {}".format(w,b))

Weights: 
 <tf.Variable 'dense/kernel:0' shape=(1, 1) dtype=float32, numpy=array([[-0.1450262]], dtype=float32)> 
 Bias: 
 <tf.Variable 'dense/bias:0' shape=(1,) dtype=float32, numpy=array([0.], dtype=float32)>

This results in one weight and one bias, not all weights generated by model at every batch /epoch.

And second method to get weights directly from h5 file:

# Functions to read weights from h5 file
import h5py

def getH5Keys(fileName):

    keys = []
    with h5py.File(fileName, mode='r') as f:
        for key in f:
            keys.append(key)

    return keys

def isGroup(obj):
    if isinstance(obj, h5py.Group):
        return True
    else:
        return False

def isDataset(obj):
    if isinstance(obj, h5py.Dataset):
        return True
    else:
        return False

def getDataSetsFromGroup(datasets, obj):
    if isGroup(obj):
        for key in obj:
            x = obj[key]
            getDataSetsFromGroup(datasets, x)
    else:
        datasets.append(obj)


def getWeightsForLayer(layerName, fileName):

    weights = []
    with h5py.File(fileName, mode='r') as f:
        for key in f:
            if layerName in key:
                obj = f[key]
                datasets = []
                getDataSetsFromGroup(datasets, obj)

                for dataset in datasets:
                    w = np.array(dataset)
                    weights.append(w)

    return weights

This method returns the same singular values for one weight and one bias:

layers = getH5Keys(checkpoint_filepath)
firstLayer = layers[0]
print(layers) # ['dense']

weights = getWeightsForLayer(firstLayer, checkpoint_filepath)
for w in weights:
    print(w.shape)
print(weights)

Output:

(1,)
(1, 1)
[array([0.], dtype=float32), array([[-0.1450262]], dtype=float32)]

Again I get only one weight and one bias. How to get all weights generated by model for every batch /epoch?

Update

Answer from 10xAI works for me. However, in my case I have one level of network with one unit, so I access weights and bias differently:

weights_dict = {}
weight_callback = tf.keras.callbacks.LambdaCallback \
( on_epoch_end=lambda epoch, logs:  weights_dict.update({epoch:model.get_weights()}))

# Train model
history = model.fit(xs, ys, batch_size=batch_size, epochs=epochs, 
                    callbacks=[weight_callback])

print(weights_dict[0])
Output: [array([[1.5375139]], dtype=float32), array([0.00499998], dtype=float32)]

print("*** Epoch: ", epoch, "\nWeight: ", weights_dict[0][0][0], " bias: ", weights_dict[1][0])
Output: *** Epoch:  79 
Weight:  [1.5375139]  bias:  [[1.5424858]]

Answer 1

You may use lambda callback and save it in a dictionary.

weights_dict = {}

weight_callback = tf.keras.callbacks.LambdaCallback \
( on_epoch_end=lambda epoch, logs: weights_dict.update({epoch:model.get_weights()}))

history = model.fit( x_train, y_train, batch_size=16, epochs=5, callbacks=weight_callback )

# retrive weights
for epoch,weights in weights_dict.items():
    print("Weights for 2nd Layer of epoch #",epoch+1)
    print(weights[2])
    print("Bias for 2nd Layer of epoch #",epoch+1)
    print(weights[3])

You can create it for batch level too.

Answer 2

The Keras Documentation indicates that filename can have the epoch number in the filename. To allow each set of data to be preserved consider using a filename more like:

checkpoint_filepath = './checkpoint-{epoch:02d}.hdf5'

filepath: string or PathLike, path to save the model file.

• filepath can contain named formatting options, which will be filled the value of epoch and keys in logs (passed in on_epoch_end). For example: if filepath is weights.{epoch:02d}-{val_loss:.2f}.hdf5, then the model checkpoints will be saved with the epoch number and the validation loss in the filename.