python - 有没有办法为给定的小批量打印每层的计算最大梯度？

Question

我正在使用 MNIST 数据集实现一个完全连接的分类模型。部分代码如下：

model=tf.keras.models.Sequential([
tf.keras.layers.Input(shape=(28, 28, 1)),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(32, activation='relu'),
tf.keras.layers.Dense(32, activation='relu'),
tf.keras.layers.Dense(32, activation='relu'),
tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(
    loss='categorical_crossentropy',
    optimizer=tf.optimizers.SGD(),
    metrics=["accuracy"]
)

model.fit(
    x_train,
    y_train,
    batch_size=64,
    epochs=3,
    validation_data=(x_test, y_test)
)

有没有办法为给定的小批量打印每一层的最大梯度？

Answer 1

定义自定义训练循环而不是调用compile()and fit()。

optimizer=tf.keras.optimizers.Adam(0.001)
loss=tf.keras.losses.SparseCategoricalCrossentropy()

for x, y in zip(x_train, y_train):
    with tf.GradientTape() as tape:
        predictions = model(x)
        loss_value = loss(y, predictions)
    gradients = tape.gradient(loss_value, model.trainable_weights)
    grads_and_vars = zip(gradients, model.trainable_weights)
    optimizer.apply_gradients(grads_and_vars)
    for layer in range(0, 4): # for 4 layers
        print('max gradient of layer={}, kernel={}, bias={}'.format(
            layer, gradients[layer].numpy().max(), gradients[layer*2+1].numpy().max()))
        

看看这个：关于 Keras