自定义GCAdam优化器get_gradients方法内print语句fit运行时无输出问题

问题原因 #

• TensorFlow 静态图执行机制限制
  TensorFlow 默认会把model.fit的训练逻辑编译为静态计算图优化执行，静态图内的普通Python print 不会被执行，输出自然无法打印到终端。__init__方法是在Python侧初始化优化器实例时执行的，处于eager执行环境，所以print('dddd')可以正常输出。
• 代码存在参数传递错误
  你调用父类get_gradients方法时没有传入要求的loss和params参数，这行代码会直接抛出异常，print('ttt')根本没有执行机会。由于异常发生在静态图构建/执行阶段，错误信息没有直接输出到终端，你没感知到报错。

解决方法 #

1. 修复参数错误，把gradients = super().get_gradients()修改为：

gradients = super().get_gradients(loss, params)

2. 调试时如果需要打印内容，将普通print替换为TensorFlow内置的tf.print：

tf.print('ttt')

tf.print是TensorFlow官方提供的打印算子，会被嵌入静态计算图执行，可以正常输出内容。
3. 如果需要用普通Python print调试，可以在编译模型时开启eager执行模式：

model.compile(loss="categorical_crossentropy", optimizer=optimizergc, metrics=["accuracy"], run_eagerly=True)

该模式下训练流程不会编译为静态图，会逐行执行Python代码，普通print可以正常输出，但训练速度会大幅下降，仅适合调试使用。

修正后的完整代码示例 #

from tensorflow.keras.optimizers import Adam
import tensorflow as tf

class GCAdam(Adam):
    def __init__(self, learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-7, amsgrad=False, name='Adam', **kwargs):
        super().__init__(learning_rate=learning_rate, beta_1=beta_1, beta_2=beta_2, epsilon=epsilon, amsgrad=amsgrad, name=name, **kwargs)
        print('dddd')
   
    def get_gradients(self, loss, params):
        grads = []
        # 修复父类方法参数传递
        gradients = super().get_gradients(loss, params)
        
        # 用tf.print替代普通print适配静态图
        tf.print('ttt')
        for grad in gradients:
            grad_len = len(grad.shape)
            if grad_len > 1:
                axis = list(range(grad_len - 1))
                grad -= tf.reduce_mean(grad, axis=axis, keepdims=True)
            grads.append(grad)

        return grads


optimizergc = GCAdam(learning_rate=1e-4)

# 调试阶段可以加run_eagerly=True，正式训练去掉即可
model.compile(loss="categorical_crossentropy", optimizer=optimizergc, metrics=["accuracy"])

model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_split=0.1,verbose=1)

内容的提问来源于stack exchange，提问作者lezi