TensorFlow自定义回调报NameError: myCallback未定义排查

问题场景 #

使用TensorFlow框架训练MNIST手写数字识别模型，预期实现逻辑如下：

• 加载MNIST数据集，对像素值做归一化处理
• 自定义myCallback回调类，当训练准确率达到99%时自动终止训练
• 搭建由Flatten层、全连接层组成的Sequential模型完成多分类任务

编写完代码运行时，在实例化回调类的步骤触发报错：

NameError: name 'myCallback' is not defined

错误根因 #

代码存在3个明显问题：

1. 重复定义了两个同名的train_mnist函数，myCallback类被错误嵌套在第一个无实际逻辑、无入参的train_mnist函数内部，属于该函数的局部变量，第二个真正执行训练逻辑的train_mnist函数无法访问到这个类定义，因此触发名称未定义错误
2. 回调类的on_epoch_end方法第一个参数大小写错误，定义时写为大写开头的Self，方法内部调用时写为小写开头的self，即使解决名称问题也会触发属性报错
3. 部分代码缩进混乱，作用域划分不符合预期

修复后完整代码 #

import os
import tensorflow as tf
from tensorflow import keras

# 获取数据集路径
current_dir = os.getcwd()
data_path = os.path.join(current_dir, "data/mnist.npz")

# 加载训练集，丢弃测试集
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data(path=data_path)
        
# 像素值归一化
x_train = x_train / 255.0

# 打印数据集维度信息
data_shape = x_train.shape
print(f"There are {data_shape[0]} examples with shape ({data_shape[1]}, {data_shape[2]})")

# 自定义早停回调类
class myCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs={}):
        if (logs.get('accuracy') is not None) and (logs.get('accuracy') > 0.99):
            print("\nReached 99% accuracy so cancelling training!") 
            self.model.stop_training = True

# 模型训练函数
def train_mnist(x_train, y_train):
    # 实例化回调
    callbacks = myCallback()
    
    # 构建Sequential模型
    model = tf.keras.models.Sequential([
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(512, activation=tf.nn.relu),
        tf.keras.layers.Dense(10, activation=tf.nn.softmax)
    ])
    
    # 编译模型
    model.compile(optimizer='adam', 
                  loss='sparse_categorical_crossentropy', 
                  metrics=['accuracy']) 
    
    # 传入回调启动训练，最多跑10轮
    history = model.fit(x_train, y_train, epochs=10, callbacks=[callbacks])
     
    return history

# 传入数据启动训练
hist = train_mnist(x_train, y_train)

关键修改说明 #

• 删除冗余的空train_mnist嵌套定义，将myCallback类移到全局作用域，保证训练函数实例化回调时可以正常访问类定义
• 统一回调类方法首参数的写法，将大写开头的Self修正为Python类方法约定的小写self，和方法内部的self.model调用匹配，避免后续属性错误
• 调整全量代码缩进，保证各逻辑块的作用域符合预期

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