tensorflow2.0 - 如何为我的张量流模型的特定层打印梯度值？

Question

我正在 TensorFlow 2.3 中训练模型

我想查看模型不同层的渐变值。

我知道如何访问模型的各个层。

with tf.GradientTape() as tape:
    lstm = model(song_emb_inp=song_emb_id_x_batch,
                 time_bucket_emb_inp=time_bucket_emb_id_x_batch,
                 training=True)
    
    loss = compute_loss(model, song_emb_id_y_batch, lstm)   
    
grads = tape.gradient(loss, model.trainable_variables)

for var, g in zip(model.trainable_variables, grads):
    print(f'{var.name}, shape: {g.shape}')     

它提供以下输出 -

如何打印所有这些可训练层的梯度值？