基于Machine-Learning-Collection项目的TensorFlow自定义训练循环教程-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_01149/article/details/148441358

基于Machine-Learning-Collection项目的TensorFlow自定义训练循环教程

Machine-Learning-Collection A resource for learning about Machine learning & Deep Learning 项目地址: https://gitcode.com/gh_mirrors/ma/Machine-Learning-Collection

前言

在深度学习模型训练过程中，TensorFlow提供了两种主要的训练方式：使用内置的fit()方法和自定义训练循环。本教程将重点介绍如何使用TensorFlow实现自定义训练循环，这是从Machine-Learning-Collection项目中提取的一个实用示例。

环境准备与数据加载

首先，我们需要设置TensorFlow环境并加载MNIST数据集：

import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # 减少TensorFlow日志输出

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.datasets import mnist
import tensorflow_datasets as tfds

# 启用GPU内存动态增长
physical_devices = tf.config.list_physical_devices("GPU")
tf.config.experimental.set_memory_growth(physical_devices[0], True)

# 加载MNIST数据集
(ds_train, ds_test), ds_info = tfds.load(
    "mnist",
    split=["train", "test"],
    shuffle_files=True,
    as_supervised=True,
    with_info=True,
)

数据预处理

良好的数据预处理是模型训练成功的关键。我们定义了一个归一化函数，并对训练集和测试集进行预处理：

def normalize_img(image, label):
    """将图像像素值归一化到0-1范围"""
    return tf.cast(image, tf.float32) / 255.0, label

AUTOTUNE = tf.data.experimental.AUTOTUNE
BATCH_SIZE = 128

# 训练集预处理流程
ds_train = ds_train.map(normalize_img, num_parallel_calls=AUTOTUNE)
ds_train = ds_train.cache()  # 缓存数据集以提高性能
ds_train = ds_train.shuffle(ds_info.splits["train"].num_examples)
ds_train = ds_train.batch(BATCH_SIZE)
ds_train = ds_train.prefetch(AUTOTUNE)  # 预取数据以减少等待时间

# 测试集预处理流程
ds_test = ds_train.map(normalize_img, num_parallel_calls=AUTOTUNE)
ds_test = ds_train.batch(128)
ds_test = ds_train.prefetch(AUTOTUNE)

模型构建

我们构建一个简单的卷积神经网络模型：

model = keras.Sequential(
    [
        keras.Input((28, 28, 1)),  # MNIST图像输入尺寸
        layers.Conv2D(32, 3, activation="relu"),  # 32个3x3卷积核
        layers.Flatten(),  # 展平层
        layers.Dense(10, activation="softmax"),  # 10个类别的输出层
    ]
)

自定义训练循环实现

自定义训练循环的核心在于手动控制训练过程，这为我们提供了更大的灵活性：

num_epochs = 5
optimizer = keras.optimizers.Adam()  # 使用Adam优化器
loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)  # 损失函数
acc_metric = keras.metrics.SparseCategoricalAccuracy()  # 准确率指标

# 训练循环
for epoch in range(num_epochs):
    print(f"\nStart of Training Epoch {epoch}")
    for batch_idx, (x_batch, y_batch) in enumerate(ds_train):
        with tf.GradientTape() as tape:
            y_pred = model(x_batch, training=True)
            loss = loss_fn(y_batch, y_pred)

        # 计算梯度并更新权重
        gradients = tape.gradient(loss, model.trainable_weights)
        optimizer.apply_gradients(zip(gradients, model.trainable_weights))
        
        # 更新准确率指标
        acc_metric.update_state(y_batch, y_pred)

    # 输出当前epoch的准确率
    train_acc = acc_metric.result()
    print(f"Accuracy over epoch {train_acc}")
    acc_metric.reset_states()  # 重置指标

测试循环

训练完成后，我们需要在测试集上评估模型性能：

# 测试循环
for batch_idx, (x_batch, y_batch) in enumerate(ds_test):
    y_pred = model(x_batch, training=True)
    acc_metric.update_state(y_batch, y_pred)

test_acc = acc_metric.result()
print(f"Accuracy over Test Set: {test_acc}")
acc_metric.reset_states()

技术要点解析

GradientTape机制：tf.GradientTape是TensorFlow提供的自动微分工具，它记录前向传播过程中的操作，然后可以计算任意可微量的梯度。
数据管道优化：使用cache()、prefetch()等方法可以显著提高数据加载效率，特别是在使用GPU训练时。
指标管理：keras.metrics提供了各种评估指标，可以方便地跟踪模型性能。
自定义灵活性：与内置的fit()方法相比，自定义循环允许我们：
- 更精细地控制训练过程
- 实现复杂的训练逻辑
- 更容易调试和监控