tf.keras：如何在训练过程中记录渐变？

Question

我正在训练tf.keras模型，并希望在训练过程中记录梯度。

更准确地说，我想在训练过程中的每次迭代中为每个层的内核计算梯度的范数，并将其写到CSV文件或类似的文件中。

我该怎么做？

我有一个部分解决方案，但被卡住了：

1. 编写一个自定义优化器继承的自定义tf.keras.optimizer。在该自定义优化器中，在计算梯度时抓取它们，然后执行所需的任何计算（在我的情况下：为每一层计算其范数）。
2. 现在呢？为了将这些渐变实际写入某些文件，我需要确保在训练期间（使用model.fit()）确实获取了我添加到自定义优化器中的ops的输出。我该怎么办？
3. 即使在培训期间评估了我添加的渐变操作，我仍然必须以某种方式将它们写入文件。我该怎么办？

或者也许有更好的方法？