Keras项目实战：自定义TensorFlow训练步骤详解

韩烨琰

于 2025-05-30 09:03:58 发布

阅读量229

点赞数 5

CC 4.0 BY-SA版权

本文链接：https://blog.youkuaiyun.com/gitblog_00567/article/details/148323860

Keras项目实战：自定义TensorFlow训练步骤详解

keras keras-team/keras: 是一个基于 Python 的深度学习库，它没有使用数据库。适合用于深度学习任务的开发和实现，特别是对于需要使用 Python 深度学习库的场景。特点是深度学习库、Python、无数据库。项目地址: https://gitcode.com/gh_mirrors/ke/keras

概述

在Keras框架中，fit()方法是训练模型的标准方式，它提供了许多便利功能，如回调函数、分布式训练支持等。然而，当我们需要实现自定义训练逻辑时，就需要深入了解如何覆盖默认的训练步骤。本文将详细介绍如何在Keras中自定义训练步骤，同时保留fit()方法的便利性。

为什么需要自定义训练步骤

Keras遵循渐进式复杂性披露原则，这意味着：

对于标准监督学习任务，直接使用fit()即可
对于需要完全控制的情况，可以编写完整的训练循环
对于介于两者之间的情况，可以只覆盖训练步骤

这种设计让开发者能够根据需求灵活选择控制级别，而不会因为需求稍微复杂就不得不放弃所有高级功能。

基础示例：覆盖train_step

让我们从一个简单的例子开始，展示如何通过继承keras.Model类并覆盖train_step方法来实现自定义训练逻辑：

class CustomModel(keras.Model):
    def train_step(self, data):
        # 解包数据
        x, y = data
        
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # 前向传播
            loss = self.compute_loss(y=y, y_pred=y_pred)  # 计算损失
        
        # 计算梯度
        gradients = tape.gradient(loss, self.trainable_variables)
        
        # 更新权重
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
        
        # 更新指标
        for metric in self.metrics:
            if metric.name == "loss":
                metric.update_state(loss)
            else:
                metric.update_state(y, y_pred)
        
        return {m.name: m.result() for m in self.metrics}

关键点说明：

data参数的结构取决于传递给fit()的内容
使用self.compute_loss()计算损失，它会自动处理编译时配置的损失函数
通过self.metrics访问和更新所有指标

更底层的实现

如果我们需要更多控制，可以完全不依赖compile()方法提供的损失和指标功能：

class CustomModel(keras.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.loss_tracker = keras.metrics.Mean(name="loss")
        self.mae_metric = keras.metrics.MeanAbsoluteError(name="mae")
        self.loss_fn = keras.losses.MeanSquaredError()
    
    def train_step(self, data):
        x, y = data
        
        # 训练步骤实现...
        
        return {"loss": self.loss_tracker.result(), "mae": self.mae_metric.result()}
    
    @property
    def metrics(self):
        return [self.loss_tracker, self.mae_metric]

在这种实现中：

我们完全自己管理指标
metrics属性确保指标在每个epoch开始时自动重置
仍然可以使用fit()方法的所有便利功能

支持样本权重

在实际应用中，我们经常需要处理样本权重或类别权重。以下是支持这些功能的实现方式：

class CustomModel(keras.Model):
    def train_step(self, data):
        if len(data) == 3:
            x, y, sample_weight = data
        else:
            x, y = data
            sample_weight = None
        
        # 计算损失时传入sample_weight
        loss = self.compute_loss(y=y, y_pred=y_pred, sample_weight=sample_weight)
        
        # 更新指标时也传入sample_weight
        metric.update_state(y, y_pred, sample_weight=sample_weight)

自定义评估步骤

与训练步骤类似，我们也可以自定义评估步骤：

class CustomModel(keras.Model):
    def test_step(self, data):
        x, y = data
        y_pred = self(x, training=False)
        loss = self.compute_loss(y=y, y_pred=y_pred)
        
        # 更新指标...
        
        return {m.name: m.result() for m in self.metrics}

完整案例：GAN实现

让我们看一个完整的生成对抗网络(GAN)实现，展示如何在一个类中管理两个模型：

class GAN(keras.Model):
    def __init__(self, discriminator, generator, latent_dim):
        super().__init__()
        self.discriminator = discriminator
        self.generator = generator
        self.latent_dim = latent_dim
        self.d_loss_tracker = keras.metrics.Mean(name="d_loss")
        self.g_loss_tracker = keras.metrics.Mean(name="g_loss")
    
    def compile(self, d_optimizer, g_optimizer, loss_fn):
        super().compile()
        self.d_optimizer = d_optimizer
        self.g_optimizer = g_optimizer
        self.loss_fn = loss_fn
    
    def train_step(self, real_images):
        # 训练判别器...
        # 训练生成器...
        
        return {
            "d_loss": self.d_loss_tracker.result(),
            "g_loss": self.g_loss_tracker.result()
        }
    
    @property
    def metrics(self):
        return [self.d_loss_tracker, self.g_loss_tracker]

这个GAN实现展示了：