具有批处理归一化层的渐变的Keras张量板抛出值错误

Question

我正在使用Tensorflow上的Keras训练卷积神经网络，并且想在Tensorboard上看到梯度的直方图。

我的回调如下：

tensorboard = TensorBoard(log_dir=log_dir, histogram_freq=1, write_grads=True)

当我拟合模型时正在使用它：

history = model.fit_generator(
    gen_train, steps_per_epoch=ceil(num_samples_train/batch_size),
    epochs=epochs, validation_data=gen_val, validation_steps=ceil(num_samples_val/batch_size), verbose=1,
    callbacks=[tensorboard, checkpoint])

但是，如果我的模型具有批处理规范化层：

model.add(BatchNormalization(axis=channels))

我收到以下错误：

  File "\.venv\lib\site-packages\keras\optimizers.py", line 91, in get_gradients
raise ValueError('An operation has `None` for gradient. '
ValueError: An operation has `None` for gradient. Please make sure that all of your ops have a gradient defined (i.e. are differentiable). Common ops without gradient: K.argmax, K.round, K.eval.

如果我的模型仅具有卷积和激活，则不会出现错误。但是我的理解是，即使网络具有批量归一化，也应绘制梯度。有什么想法为什么这行不通吗？如果实际上这是预期的行为，那么是否有一种变通方法来在Tensorboard上获取渐变？