keras TensorBoard write_grads prevents training from starting - Python

I am trying to use the TensorBoard callback to visualize the training of my network. The network reuses parts of ResNet (pre-trained and untrainable for my purposes). When omitting the write_grads=True option everything works as expected and I can see the histograms for my own trainable layers in TensorBoard. However with write_grads set to true the last visible output is Train on 30 samples, validate on 10 samples and the training itself never seems to start.

I noticed that the main memory used by the python process slowly increases from a couple of hundred megabytes to several gigabytes during that time. It came to my mind that saving all gradients in ResNet might take a significant amount of space. But since most of it is untrainable in my particular case I think this is not a general problem but rather an implementation issue.

Please find below a script to reproduce the issue.

import numpy as np

from keras.callbacks import TensorBoard
from keras.models import Model
from keras.layers import Input, Dense, Conv2D, GlobalMaxPool2D
from keras.applications.resnet50 import ResNet50

x = np.random.normal(size=(30, 200, 200, 3))
xv = np.random.normal(size=(10, 200, 200, 3))
y = np.random.randint(0, 2, size=(30, 5))
yv = np.random.randint(0, 2, size=(10, 5))

tb = TensorBoard(log_dir=r'C:\logs', histogram_freq=1, batch_size=2, write_grads=True)

# Build network that reuses a pre-trained part of ResNet and is then followed by some other layers
# Note that the ResNet part is untrainable!
inputs = Input((200, 200, 3))
net = ResNet50(input_tensor=inputs, weights='imagenet', include_top=False)
for layer in net.layers:
    layer.trainable = False
net = Conv2D(10, 5, activation='relu')(net.layers[141].output)
net = GlobalMaxPool2D()(net)
net = Dense(5, activation='softmax')(net)

model = Model(inputs=inputs, outputs=net)
model.compile(optimizer='sgd', loss='binary_crossentropy')
model.fit(x, y, batch_size=2, validation_data=(xv, yv), epochs=10, callbacks=[tb])

• [X] Check that you are up-to-date with the master branch of Keras.
• [X] If running on TensorFlow, check that you are up-to-date with the latest version.

Asked Oct 17 '21 12:10

jprellberg