如何在Google Colab中打印TensorFlow分类器训练的损失与步数

我正在跟随FreeCodeCamp机器学习课程，使用以下TensorFlow代码训练鸢尾花分类数据集：

CSV_COLUMN_NAMES = ['SepalLength', 'SepalWidth', 'PetalLength', 'PetalWidth', 'Species']
SPECIES = ['Setosa', 'Versicolor', 'Virginica'] # 定义后续要用的常量

train_path = tf.keras.utils.get_file(
    "iris_training.csv", 
    "https://storage.googleapis.com/download.tensorflow.org/data/iris_training.csv")
test_path = tf.keras.utils.get_file(
    "iris_test.csv", 
    "https://storage.googleapis.com/download.tensorflow.org/data/iris_test.csv")

# 用Keras（TensorFlow内置模块）获取数据集并读取为pandas DataFrame
train = pd.read_csv(train_path, names=CSV_COLUMN_NAMES, header=0)
test = pd.read_csv(test_path, names=CSV_COLUMN_NAMES, header=0) 

train.rename(columns={'120': 'SepalLength', '4': 'SepalWidth',
                      'setosa': 'PetalLength', 'versicolor': 'PetalWidth',
                      'virginica': 'Species'}, inplace=True)
test.rename(columns={'30': 'SepalLength', '4': 'SepalWidth',
                     'setosa': 'PetalLength', 'versicolor': 'PetalWidth',
                     'virginica': 'Species'}, inplace=True)

train_y = train.pop('Species')
test_y = test.pop('Species')
train.head() # 此时Species列已被移除

def input_fn(features, labels, training=True, batch_size=256):
  # 将输入转换为Dataset
  dataset = tf.data.Dataset.from_tensor_slices((dict(features), labels))

  # 训练模式下打乱并重复数据
  if training:
    dataset = dataset.shuffle(1000).repeat()
  return dataset.batch(batch_size)

# 特征列定义输入的使用方式
my_feature_columns = [] 
for key in train.keys():
  my_feature_columns.append(tf.feature_column.numeric_column(key=key))

print(my_feature_columns) 

# 构建包含2个隐藏层（30和10个节点）的DNN模型
classifier = tf.estimator.DNNClassifier(
  feature_columns=my_feature_columns,  
  hidden_units=[30, 10], # 两个隐藏层分别有30和10个节点
  n_classes=3) # 模型需要区分3个类别

classifier.train(
    input_fn=lambda: input_fn(train, train_y, training=True),
    steps=5000) # 定义训练的总步数

# TODO: 打印损失和步数

本地运行时classifier.train()会在控制台输出训练步数和损失值，但在Google Colab中无此输出，请问如何在Colab中打印训练过程中的损失与步数？

解决办法 #

方法一：调整TensorFlow日志级别 #

Colab默认调低了TensorFlow的日志输出级别，只需在代码开头添加日志配置，即可让训练过程的信息正常显示：

import tensorflow as tf
tf.get_logger().setLevel('INFO')

添加后，Estimator训练时的步数、损失等日志会直接输出到Colab控制台。

方法二：自定义Hook捕获训练指标 #

如果需要自定义输出频率或格式，可以通过SessionRunHook来捕获训练过程中的损失和步数：

import tensorflow as tf

class TrainingLoggerHook(tf.estimator.SessionRunHook):
    def begin(self):
        self.step = 0

    def before_run(self, run_context):
        # 捕获总损失值
        return tf.estimator.SessionRunArgs(tf.compat.v1.losses.get_total_loss())

    def after_run(self, run_context, run_values):
        self.step += 1
        # 每500步打印一次，可以根据需求修改间隔
        if self.step % 500 == 0:
            loss_value = run_values.results
            print(f"训练步数: {self.step}, 当前损失: {loss_value:.4f}")

# 训练时传入自定义Hook
classifier.train(
    input_fn=lambda: input_fn(train, train_y, training=True),
    steps=5000,
    hooks=[TrainingLoggerHook()]
)

方法三：切换到Keras API（推荐） #

TensorFlow Estimator API相对老旧，Keras API在Colab中的日志输出更直观，使用也更简洁。改用Keras实现后，训练过程会自动输出步数和损失：

# 构建Keras序列模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(30, activation='relu', input_shape=(4,)),
    tf.keras.layers.Dense(10, activation='relu'),
    tf.keras.layers.Dense(3, activation='softmax')
])

# 编译模型
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# 训练模型，自动输出每轮的损失和准确率
history = model.fit(train, train_y, epochs=50, batch_size=256)

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