企业官网建设流程全解析

1. 多损失函数这件事，为什么值得单独拎出来聊

做深度学习项目的人，迟早会撞上这样一个场景：模型要同时预测多个目标，或者一个目标需要从不同维度去约束。比如一个多任务网络，既要预测点击率，又要预测转化率，还要预测停留时长；又比如一个图像生成模型，既要保证像素级重建精度，又要保证感知质量，还得让风格分布对齐。这时候，单个损失函数根本不够用，你得把多个损失加权组合起来。

听起来很简单对吧？total_loss = a * loss1 + b * loss2 + c * loss3，一行代码的事。但真正动起手来，问题一个接一个冒出来：自定义损失函数怎么传参？多个损失函数怎么分别监控？不同损失的量级差异太大怎么办？训练过程中想动态调整权重又该怎么写？更别提还要兼顾model.compile()和model.fit()的流程，以及自定义训练循环里的灵活控制。

我用 Keras 做多任务学习断断续续有几年了，踩过的坑从“损失函数不收敛”到“梯度爆炸但找不到原因”，再到“验证集损失和训练集损失计算方式不一致导致模型选择错误”，几乎每个环节都翻过车。这篇文章就把这些经验系统性地梳理一遍，从设计思路到代码实现，从静态权重到动态调整，从compile/fit到自定义训练循环，把“优雅实现多个自定义损失函数”这件事讲透。

适合谁看？如果你已经能用 Keras 搭个简单模型跑通训练，现在需要处理多目标优化、多任务学习、或者自定义复杂的损失组合，那这篇内容就是为你准备的。我会尽量用从业者之间交流的方式来讲，少堆术语，多讲“为什么这么设计”和“实际跑起来会怎样”。

2. 整体设计思路：多损失函数到底该怎么组织

2.1 先想清楚：多损失的本质是什么

多损失函数的本质，是在优化过程中同时约束模型的多个行为。每个损失函数代表一种“期望”，模型参数需要在这些期望之间找到平衡点。这个平衡点怎么找，取决于你怎么组合这些损失。

最朴素的方式是加权求和：

total_loss = w1 * loss1 + w2 * loss2 + w3 * loss3

但加权求和有个隐含假设：所有损失函数对模型参数的影响是可加的，且权重是固定的。实际项目中，这个假设经常不成立。比如一个损失在训练初期梯度很大，另一个损失在训练后期才起作用，固定权重就很难让两者都发挥好。

所以设计多损失方案时，核心要回答三个问题：

• 损失之间是否量级匹配？如果一个损失是 MSE（量级可能在 0.01），另一个是交叉熵（量级可能在 2.0），直接相加会导致量级大的主导梯度。
• 权重是静态还是动态？静态权重简单但不够灵活，动态权重（如不确定性加权、GradNorm）更智能但实现复杂。
• 监控和调试怎么做？如果只看到 total_loss，根本不知道哪个子损失在起作用、哪个在拖后腿。

2.2 三种主流组织方式对比

在实际项目中，我常用三种方式来组织多损失函数，各有适用场景：

我个人的经验是：如果模型本身就是多输出的，优先用第二种方式，因为 Keras 的compile原生支持loss={'output1': loss1, 'output2': loss2}这种写法，而且fit会自动记录每个输出的损失值，调试起来非常方便。如果模型是单输出但需要多个损失约束（比如特征层的正则化、中间层的辅助损失），那就得用第一种或第三种。

2.3 为什么“优雅”很重要

“优雅”在这里不是指代码好看，而是指可维护、可调试、可扩展。我见过太多项目，多损失函数写成一坨，权重硬编码在 loss 函数里，想调个权重得改代码重新跑，训练日志里只有一个 total_loss，模型不收敛时完全不知道从哪下手。

优雅的实现应该满足：

• 损失函数定义清晰，每个损失独立可测试
• 权重可配置，最好能通过配置文件或命令行参数调整
• 训练过程中能分别看到每个子损失的值
• 支持动态权重策略，不需要改模型结构
• 兼容compile/fit和自定义训练循环两种模式

下面我就按这个标准，一步步拆解具体实现。

3. 核心细节解析：自定义损失函数的传参与组合

3.1 自定义损失函数的基本写法

Keras 里自定义损失函数有两种写法：函数式和类式。函数式简单直接：

import tensorflow as tf from tensorflow import keras def custom_mse(y_true, y_pred): return tf.reduce_mean(tf.square(y_true - y_pred))

但函数式有个致命问题：没法传额外参数。比如你想给损失函数加一个scale参数，或者传入一个 mask，函数式就做不到了。这时候需要用类式：

class WeightedMSE(keras.losses.Loss): def __init__(self, scale=1.0, name='weighted_mse'): super().__init__(name=name) self.scale = scale def call(self, y_true, y_pred): return self.scale * tf.reduce_mean(tf.square(y_true - y_pred))

类式的好处是参数在__init__里传，call里直接用。而且继承keras.losses.Loss后，Keras 会自动处理一些边界情况，比如损失函数的 reduction 方式。

注意：如果你用的是tf.keras，继承keras.losses.Loss时，call方法的返回值应该是每个样本的损失（shape 为[batch_size]），而不是已经 reduce 过的标量。Keras 会在外部做 reduction。但如果你直接传函数式损失给compile，Keras 会默认做reduce_mean。这个差异很容易搞混，我建议统一用类式，行为更可控。

3.2 多个损失函数的组合方式

假设你有三个损失：loss_a、loss_b、loss_c，权重分别是0.5、0.3、0.2。最直接的组合方式是在一个函数里算完：

def combined_loss(y_true, y_pred): a = loss_a(y_true, y_pred) b = loss_b(y_true, y_pred) c = loss_c(y_true, y_pred) return 0.5 * a + 0.3 * b + 0.2 * c

但这样写有个问题：y_true和y_pred的结构必须同时满足三个损失的要求。如果loss_a需要的是分类标签，loss_b需要的是回归目标，那y_true就得是一个包含多个字段的结构，y_pred也得对应多个输出。这种情况下，更清晰的方式是让模型多输出，每个输出对应一个损失。

# 模型定义 input_layer = keras.Input(shape=(128,)) shared = keras.layers.Dense(64, activation='relu')(input_layer) output_a = keras.layers.Dense(10, activation='softmax', name='output_a')(shared) output_b = keras.layers.Dense(1, name='output_b')(shared) output_c = keras.layers.Dense(1, name='output_c')(shared) model = keras.Model(inputs=input_layer, outputs=[output_a, output_b, output_c]) # 编译时指定每个输出的损失 model.compile( optimizer='adam', loss={ 'output_a': keras.losses.CategoricalCrossentropy(), 'output_b': WeightedMSE(scale=1.0), 'output_c': WeightedMSE(scale=0.5) }, loss_weights={'output_a': 0.5, 'output_b': 0.3, 'output_c': 0.2}, metrics={'output_a': 'accuracy'} )

这种写法非常清晰：每个输出的损失独立定义，loss_weights控制组合权重，metrics可以单独监控某个输出的准确率。训练时fit会返回一个 history 字典，里面包含output_a_loss、output_b_loss、output_c_loss和总的loss，调试起来一目了然。

3.3 损失量级不匹配怎么办

这是多损失函数最常见的坑。假设loss_a是交叉熵，初始值在 2.3 左右（10 分类随机猜测），loss_b是 MSE，初始值在 0.01 左右。如果权重都是 1.0，那 total_loss 几乎完全由loss_a主导，loss_b的梯度贡献可以忽略不计，模型根本不会去优化loss_b对应的目标。

解决办法有三种：

第一种：手动调权重。根据经验或实验，给量级小的损失更大的权重。比如loss_b权重设为 100，让它的贡献和loss_a相当。但这个方法很粗糙，因为损失量级在训练过程中会变化，固定权重很难一直保持平衡。

第二种：归一化损失。在组合之前，把每个损失除以其初始值或滑动平均值，让它们都在相近的量级。比如：

class NormalizedLoss(keras.losses.Loss): def __init__(self, base_loss, name='normalized_loss'): super().__init__(name=name) self.base_loss = base_loss self.running_mean = tf.Variable(1.0, trainable=False) def call(self, y_true, y_pred): loss = self.base_loss(y_true, y_pred) self.running_mean.assign(0.99 * self.running_mean + 0.01 * tf.reduce_mean(loss)) return loss / (self.running_mean + 1e-8)

这个方法的好处是自适应，不需要手动调权重。但要注意，running_mean的更新应该在训练过程中进行，验证时不应该更新。实现时需要用tf.cond或training标志来控制。

第三种：不确定性加权。这是 Kendall 等人在 2018 年提出的一种方法，把每个损失的权重也作为可学习参数，通过最大化高斯似然来优化。简单来说，每个损失对应一个可学习的log_var，最终的损失是：

total_loss = sum(0.5 * exp(-log_var_i) * loss_i + 0.5 * log_var_i)

这个方法在理论上很优雅，实际效果也不错，但实现起来稍微复杂一点，需要把log_var作为模型的可训练变量。我一般会在研究性项目里用，生产环境还是倾向于手动调权重加归一化。

3.4 动态权重调整的几种策略

除了不确定性加权，还有几种常见的动态权重策略：

• GradNorm：通过平衡不同任务梯度的大小来调整权重，让每个任务的训练速度相近。
• PCGrad：当不同任务的梯度方向冲突时，把冲突的梯度投影到正交方向，减少负迁移。
• 动态平均：根据每个损失的变化率来调整权重，变化慢的损失加大权重。

这些方法各有论文支撑，实现起来都需要自定义训练循环。我的建议是：先从简单的手动权重加归一化开始，如果效果不理想再考虑复杂策略。很多时候，问题不在于权重策略不够高级，而在于模型结构或数据本身有问题。

4. 实操过程：从零搭建一个多损失训练流程

4.1 场景设定与数据准备

假设我们要做一个多任务模型，输入是用户行为特征（128 维），输出三个目标：

• 是否点击（二分类）
• 停留时长（回归）
• 是否分享（二分类）

三个任务共享底层特征，但各自有独立的输出头。数据是模拟生成的，每个样本包含 128 维特征和三个标签。

import numpy as np np.random.seed(42) n_samples = 10000 n_features = 128 X = np.random.randn(n_samples, n_features).astype(np.float32) y_click = np.random.randint(0, 2, size=(n_samples, 1)).astype(np.float32) y_duration = np.random.exponential(scale=5.0, size=(n_samples, 1)).astype(np.float32) y_share = np.random.randint(0, 2, size=(n_samples, 1)).astype(np.float32) # 划分训练集和验证集 split = int(0.8 * n_samples) X_train, X_val = X[:split], X[split:] y_click_train, y_click_val = y_click[:split], y_click[split:] y_duration_train, y_duration_val = y_duration[:split], y_duration[split:] y_share_train, y_share_val = y_share[:split], y_share[split:]

4.2 模型结构定义

模型结构很直接：共享层 + 三个独立输出头。

def build_multi_task_model(n_features): input_layer = keras.Input(shape=(n_features,), name='input') # 共享层 shared = keras.layers.Dense(256, activation='relu')(input_layer) shared = keras.layers.Dropout(0.3)(shared) shared = keras.layers.Dense(128, activation='relu')(shared) shared = keras.layers.Dropout(0.3)(shared) # 任务一：点击预测 click_head = keras.layers.Dense(64, activation='relu')(shared) click_output = keras.layers.Dense(1, activation='sigmoid', name='click')(click_head) # 任务二：停留时长预测 duration_head = keras.layers.Dense(64, activation='relu')(shared) duration_output = keras.layers.Dense(1, activation='linear', name='duration')(duration_head) # 任务三：分享预测 share_head = keras.layers.Dense(64, activation='relu')(shared) share_output = keras.layers.Dense(1, activation='sigmoid', name='share')(share_head) model = keras.Model( inputs=input_layer, outputs=[click_output, duration_output, share_output] ) return model model = build_multi_task_model(n_features) model.summary()

这里有个细节：三个输出头的最后一层激活函数不同。点击和分享是二分类，用sigmoid；停留时长是回归，用linear。这个区别很重要，因为损失函数的选择依赖于输出类型。

4.3 自定义损失函数的实现

对于二分类任务，用BinaryCrossentropy就够了。但停留时长这个回归任务，我想用 Huber 损失，因为它对异常值更鲁棒。而且我想给 Huber 损失加一个可配置的delta参数。

class HuberLoss(keras.losses.Loss): def __init__(self, delta=1.0, name='huber_loss'): super().__init__(name=name) self.delta = delta def call(self, y_true, y_pred): error = y_true - y_pred abs_error = tf.abs(error) quadratic = tf.minimum(abs_error, self.delta) linear = abs_error - quadratic return 0.5 * tf.square(quadratic) + self.delta * linear

另外，我想给点击和分享任务加一个focal loss的变体，因为实际场景中正样本可能很少，普通交叉熵会被负样本主导。Focal loss 通过调制因子让模型更关注难分类的样本。

class BinaryFocalLoss(keras.losses.Loss): def __init__(self, alpha=0.25, gamma=2.0, name='binary_focal_loss'): super().__init__(name=name) self.alpha = alpha self.gamma = gamma def call(self, y_true, y_pred): y_pred = tf.clip_by_value(y_pred, 1e-7, 1 - 1e-7) bce = -y_true * tf.math.log(y_pred) - (1 - y_true) * tf.math.log(1 - y_pred) p_t = y_true * y_pred + (1 - y_true) * (1 - y_pred) alpha_t = y_true * self.alpha + (1 - y_true) * (1 - self.alpha) focal = alpha_t * tf.pow(1 - p_t, self.gamma) * bce return focal

注意：Focal loss 的alpha和gamma需要根据数据分布来调。如果正样本比例在 10% 左右，alpha=0.25是个不错的起点。gamma一般取 2.0，增大gamma会让模型更关注难样本，但太大容易导致训练不稳定。

4.4 编译与训练配置

现在把损失函数组合起来。点击和分享用 Focal Loss，停留时长用 Huber Loss。权重方面，我先给点击 1.0，停留时长 0.5，分享 0.8。这个权重不是拍脑袋定的，而是根据验证集上的表现逐步调的。

model.compile( optimizer=keras.optimizers.Adam(learning_rate=1e-3), loss={ 'click': BinaryFocalLoss(alpha=0.25, gamma=2.0), 'duration': HuberLoss(delta=1.0), 'share': BinaryFocalLoss(alpha=0.3, gamma=1.5) }, loss_weights={ 'click': 1.0, 'duration': 0.5, 'share': 0.8 }, metrics={ 'click': ['accuracy', keras.metrics.AUC(name='auc')], 'duration': [keras.metrics.MeanAbsoluteError(name='mae')], 'share': ['accuracy', keras.metrics.AUC(name='auc')] } )

训练时用EarlyStopping和ReduceLROnPlateau来防止过拟合和动态调学习率。

callbacks = [ keras.callbacks.EarlyStopping( monitor='val_loss', patience=5, restore_best_weights=True ), keras.callbacks.ReduceLROnPlateau( monitor='val_loss', factor=0.5, patience=3, min_lr=1e-6 ) ] history = model.fit( X_train, {'click': y_click_train, 'duration': y_duration_train, 'share': y_share_train}, validation_data=( X_val, {'click': y_click_val, 'duration': y_duration_val, 'share': y_share_val} ), batch_size=64, epochs=50, callbacks=callbacks, verbose=1 )

4.5 训练过程监控与可视化

fit返回的history里包含了每个输出的损失和总损失。我习惯把关键指标画出来，看看每个任务的学习曲线是否正常。

import matplotlib.pyplot as plt fig, axes = plt.subplots(2, 2, figsize=(12, 10)) # 总损失 axes[0, 0].plot(history.history['loss'], label='train') axes[0, 0].plot(history.history['val_loss'], label='val') axes[0, 0].set_title('Total Loss') axes[0, 0].legend() # 点击任务损失 axes[0, 1].plot(history.history['click_loss'], label='train') axes[0, 1].plot(history.history['val_click_loss'], label='val') axes[0, 1].set_title('Click Loss') axes[0, 1].legend() # 停留时长任务损失 axes[1, 0].plot(history.history['duration_loss'], label='train') axes[1, 0].plot(history.history['val_duration_loss'], label='val') axes[1, 0].set_title('Duration Loss') axes[1, 0].legend() # 分享任务损失 axes[1, 1].plot(history.history['share_loss'], label='train') axes[1, 1].plot(history.history['val_share_loss'], label='val') axes[1, 1].set_title('Share Loss') axes[1, 1].legend() plt.tight_layout() plt.show()

通过这个图，你能快速判断哪个任务学得好、哪个任务过拟合了、哪个任务还没收敛。如果某个任务的验证损失一直在上升，说明这个任务的权重可能太大了，或者模型在这个任务上过拟合了。

5. 进阶技巧：自定义训练循环中的多损失控制

5.1 什么时候需要自定义训练循环

compile/fit能覆盖大部分场景，但有些需求它搞不定：

• 需要在每个 batch 后动态调整损失权重
• 需要对不同任务使用不同的优化器
• 需要在训练过程中对梯度做特殊处理（如梯度裁剪、梯度归一化）
• 需要实现复杂的训练策略（如对抗训练、元学习）

这时候就得自己写训练循环。Keras 提供了train_step和test_step的覆写机制，用起来比从头写 PyTorch 循环简单一些。

5.2 自定义训练循环的骨架

class MultiTaskModel(keras.Model): def __init__(self, n_features, **kwargs): super().__init__(**kwargs) self.shared = keras.layers.Dense(256, activation='relu') self.dropout1 = keras.layers.Dropout(0.3) self.shared2 = keras.layers.Dense(128, activation='relu') self.dropout2 = keras.layers.Dropout(0.3) self.click_head = keras.layers.Dense(64, activation='relu') self.click_output = keras.layers.Dense(1, activation='sigmoid') self.duration_head = keras.layers.Dense(64, activation='relu') self.duration_output = keras.layers.Dense(1, activation='linear') self.share_head = keras.layers.Dense(64, activation='relu') self.share_output = keras.layers.Dense(1, activation='sigmoid') # 可学习的损失权重 self.log_vars = tf.Variable([0.0, 0.0, 0.0], trainable=True, name='log_vars') def call(self, inputs, training=False): x = self.shared(inputs) x = self.dropout1(x, training=training) x = self.shared2(x) x = self.dropout2(x, training=training) click = self.click_output(self.click_head(x)) duration = self.duration_output(self.duration_head(x)) share = self.share_output(self.share_head(x)) return click, duration, share def compute_losses(self, y_true, y_pred): click_loss = BinaryFocalLoss(alpha=0.25, gamma=2.0)( y_true['click'], y_pred[0] ) duration_loss = HuberLoss(delta=1.0)( y_true['duration'], y_pred[1] ) share_loss = BinaryFocalLoss(alpha=0.3, gamma=1.5)( y_true['share'], y_pred[2] ) return click_loss, duration_loss, share_loss def train_step(self, data): X, y = data with tf.GradientTape() as tape: y_pred = self(X, training=True) click_loss, duration_loss, share_loss = self.compute_losses(y, y_pred) # 不确定性加权 losses = [click_loss, duration_loss, share_loss] total_loss = 0.0 for i, loss in enumerate(losses): precision = tf.exp(-self.log_vars[i]) total_loss += 0.5 * precision * loss + 0.5 * self.log_vars[i] trainable_vars = self.trainable_variables gradients = tape.gradient(total_loss, trainable_vars) self.optimizer.apply_gradients(zip(gradients, trainable_vars)) return { 'loss': total_loss, 'click_loss': click_loss, 'duration_loss': duration_loss, 'share_loss': share_loss, 'log_var_click': self.log_vars[0], 'log_var_duration': self.log_vars[1], 'log_var_share': self.log_vars[2] }

这个训练循环里，log_vars是可学习的，模型会自动调整每个损失的权重。训练初期，log_vars会倾向于增大（降低对应损失的权重），因为模型还没学好，大权重会导致梯度不稳定。随着训练进行，log_vars会逐渐找到合适的值。

5.3 动态权重调整的实操细节

不确定性加权虽然优雅，但有几个实操细节要注意：

第一，log_vars的初始化。一般初始化为 0，对应权重为 1.0。如果某个损失量级特别大，可以初始化为负值，让初始权重小一些。

第二，log_vars的学习率。如果和模型参数用同一个学习率，log_vars可能会变化太快，导致训练不稳定。我一般会给log_vars单独设置一个较小的学习率，或者用tf.clip_by_value限制它的范围。

第三，验证时的处理。验证时不应该更新log_vars，也不应该用验证损失去调log_vars。验证损失的计算应该用训练时学到的log_vars值，或者直接用固定权重。

def test_step(self, data): X, y = data y_pred = self(X, training=False) click_loss, duration_loss, share_loss = self.compute_losses(y, y_pred) # 验证时用固定的权重，不更新 log_vars total_loss = click_loss + 0.5 * duration_loss + 0.8 * share_loss return { 'loss': total_loss, 'click_loss': click_loss, 'duration_loss': duration_loss, 'share_loss': share_loss }

5.4 梯度冲突与PCGrad的简化实现

多任务学习中，不同任务的梯度方向可能冲突。比如任务 A 希望增大某个参数，任务 B 希望减小同一个参数，两个梯度相加后可能互相抵消，导致训练停滞。PCGrad 的思路是：如果两个梯度方向夹角大于 90 度（余弦相似度为负），就把其中一个梯度投影到另一个梯度的正交方向上。

def project_gradient(grad, other_grad): """把 grad 投影到 other_grad 的正交方向""" dot = tf.reduce_sum(grad * other_grad) norm_sq = tf.reduce_sum(other_grad * other_grad) + 1e-8 if dot < 0: return grad - (dot / norm_sq) * other_grad return grad

在实际训练循环里，需要分别计算每个任务的梯度，然后两两做投影，最后再求和。这个实现起来比较繁琐，而且计算量会随任务数量增加而快速增长。我的建议是：如果任务数量少于 5 个，可以试试 PCGrad；如果任务很多，还是用不确定性加权更实际。

6. 常见问题与排查技巧实录

6.1 损失不下降或震荡严重

这是最常见的问题。可能的原因和排查思路：

我遇到过一次特别隐蔽的情况：总损失在下降，但其中一个任务的准确率一直在随机水平。后来发现是那个任务的输出头初始化有问题，最后一层的 bias 初始值太大，导致 sigmoid 输出饱和，梯度几乎为零。把 bias 初始化为 0 后问题就解决了。

6.2 验证损失和训练损失计算方式不一致

这个问题很隐蔽，但影响很大。Keras 在训练时会用training=True调用模型，Dropout 和 BatchNormalization 会生效；验证时用training=False，这些层会切换到推理模式。如果你在自定义损失函数里也用了training标志，但没正确处理，就会导致训练和验证的损失计算方式不一致。

注意：自定义损失函数的call方法默认没有training参数。如果你需要在损失函数里区分训练和验证，得通过其他方式传递，比如在train_step和test_step里分别调用不同的损失计算逻辑。

6.3 多输出模型的数据格式问题

Keras 的多输出模型对数据格式很敏感。fit时y可以传字典、列表或元组，但必须和模型输出的顺序或名称对应。我建议统一用字典，因为字典不依赖顺序，而且可读性好。

# 推荐：用字典 y_train = {'click': y_click_train, 'duration': y_duration_train, 'share': y_share_train} # 也可以用列表，但必须和模型 outputs 的顺序一致 y_train = [y_click_train, y_duration_train, y_share_train]

如果模型输出是列表，但y传了字典，Keras 会尝试按名称匹配，匹配不上就报错。反过来，如果模型输出是字典，y传列表，Keras 会按顺序匹配，但容易搞错顺序。所以统一用字典是最稳妥的。

6.4 损失权重调整的经验法则

调权重是个经验活，但有几个法则可以参考：

• 初始权重让各损失的贡献相近。先跑几个 batch，看每个损失的初始值，然后按比例设权重，让w_i * loss_i大致相等。
• 权重不要频繁大幅调整。每次调整后至少跑几个 epoch 看效果，频繁调整会让训练不稳定。
• 验证集表现是最终标准。如果某个任务的验证指标一直上不去，可以适当增大它的权重，但要观察其他任务是否受影响。
• 任务重要性不同时，权重应该有差异。如果点击预测是核心任务，停留时长是辅助任务，那点击的权重应该明显大于停留时长。

6.5 常见问题速查表

7. 我个人的一些实操体会

多损失函数的实现，技术难度其实不高，难的是调试和平衡。我做过一个项目，三个任务的损失权重调了两周才找到合适的组合。后来总结出一个经验：先把每个任务单独训到收敛，记录下各自的损失量级和收敛速度，然后再组合起来调权重。这样比一上来就多任务一起训要高效得多。

另外，tf.function对自定义训练循环的性能影响很大。如果train_step里有很多 Python 逻辑（比如 if-else 判断），tf.function会反复追踪计算图，导致训练变慢。我的做法是把所有条件逻辑尽量用tf.cond或tf.where实现，保持train_step的计算图静态化。

最后分享一个小技巧：如果某个任务的损失在训练后期一直不下降，可以试试降低它的权重，而不是增大。因为后期模型已经学得差不多了，大权重反而会让梯度噪声主导更新，降低权重能让模型更稳定地收敛。这个反直觉的操作我在多个项目里验证过，效果不错。