[Python] Tensorboard histogram display of layer weights was working in tensorflow 2.12 but is...

The below code (a minimal 'toy' example using data generators for both training and validation data) works fine in 2.12 and produces the desired bias and kernel weight histograms in tensorboard for both hidden layers as well as the output layer. For 2.16 and later I just get the output layer bias and kernel displayed. I am guessing that this is due to the switch in tensorflow 2.16 to using Keras version 3, but I am unable to figure out how to fix my code so it generates all the histograms. Online suggestions about explicitly writing out the weights using tf.summary.histogram don't work, and neither does switching to using the validation data directly rather than using a generator.

import numpy as np
import tensorflow as tf
from keras.utils import Sequence
from keras.models import Sequential
from keras.layers import Input
from keras.layers import Dense
import datetime

class DataGenerator(Sequence):
def __init__(self, batch_size, data_size, **kwargs):
super().__init__(**kwargs)
self.data_size = data_size
self.batch_size = batch_size

self.data = np.random.random((data_size, features+1))
self.total_samples=self.data.shape[0]
self.indexes = np.arange(self.total_samples)

def __len__(self):
return int(np.ceil(self.total_samples / self.batch_size))

def __getitem__(self, index):
batch_indexes = self.indexes[index * self.batch_sizeindex + 1) * self.batch_size]
batch_data = self.data[batch_indexes]
return self.__data_generation(batch_data)

def __data_generation(self, batch_data):
x = batch_data[:, :-1]
y = batch_data[:, -1]
return x, y


if __name__ == '__main__':
features = 3
batch_size = 2

print(f"{tf.__version__=}")

model = Sequential([
Dense(20, name="twenty_unit_hidden"),
Dense(10, name="ten_unit_hidden"),
Dense(1, name="output")
])

data_gen = DataGenerator(batch_size, 10)
val_gen = DataGenerator(batch_size, 2)

log_dir = f"logs/fit/{tf.__version__}_" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1, write_graph=True, write_images=True, update_freq="epoch")

loss = tf.keras.losses.BinaryCrossentropy()
model.compile(loss=loss)

model.fit(data_gen, epochs=4, batch_size=batch_size, callbacks=[tensorboard_callback], validation_data=val_gen)

Continue reading...