1. 多损失函数这件事到底难在哪做过深度学习项目的人大概率都遇到过这种场景模型要同时预测多个目标比如一个网络既要做分类又要做回归或者一个序列模型既要输出主任务结果又要兼顾辅助任务的正则化。这时候单靠一个损失函数根本搞不定你得把多个损失按某种策略组合起来。听起来简单写起来一堆坑。我最早接触这个需求是在一个多任务学习的项目里当时模型需要同时优化三个目标主分类损失、一个辅助的对比损失、还有一个正则项。最开始我的做法特别原始直接在训练循环里手写tape.gradient把三个损失加起来求梯度。代码能跑但问题是训练逻辑和模型定义完全耦合在一起换一个损失组合就要改一遍训练代码调试的时候想单独看某个损失的收敛情况还得加一堆print。后来换成 Keras 的model.compilemodel.fit模式才发现多损失函数的组织方式其实有很多讲究。这篇文章想聊的就是在 Keras 里怎么优雅地实现多个自定义损失函数。所谓“优雅”我的标准是三条——代码可读、可维护、可扩展。不是那种“能跑就行”的写法而是你过两个月回头看还能秒懂、同事接手不用问你、加一个新损失不用动训练框架的写法。适合谁看如果你已经会用 Keras 搭模型、写过简单的自定义损失但一到多损失组合就有点乱那这篇正好。如果你还没写过自定义损失也没关系我会从基础讲起把关键概念都铺开。注意本文基于 TensorFlow 2.x 的 Keras API 展开如果你用的是独立 Keras 或其他后端部分 API 细节可能有差异但核心思路是通用的。2. 先搞清楚 Keras 里损失函数的几种写法在聊多损失之前得先把单个损失函数的几种实现方式理清楚。因为多损失本质上就是多个单损失的组合单损失写不对组合起来只会更乱。2.1 直接用内置损失函数Keras 内置了一大批常用损失比如CategoricalCrossentropy、BinaryCrossentropy、MeanSquaredError等等。用法很简单编译的时候传字符串或者传类实例都行model.compile(optimizeradam, losscategorical_crossentropy)或者更灵活一点model.compile( optimizeradam, losstf.keras.losses.CategoricalCrossentropy(from_logitsTrue) )内置损失的好处是经过充分测试、数值稳定、支持sample_weight等高级特性。能用内置的就别自己写这是我一贯的原则。很多人一上来就手写交叉熵结果忘了from_logits参数导致数值不稳定训练半天不收敛排查半天才发现是损失函数的问题。2.2 用函数式写法自定义损失当内置损失满足不了需求时就得自己写。最简单的方式是写一个函数接收y_true和y_pred返回一个标量损失def custom_mse(y_true, y_pred): diff y_true - y_pred return tf.reduce_mean(tf.square(diff))这种写法适合逻辑简单的损失。但有个坑要注意函数里必须用tf.*的算子不能用 NumPy 或 Python 原生运算否则在图模式下会出问题。我见过有人写np.mean(...)在 eager 模式下能跑一开tf.function就报错。2.3 用类继承方式自定义损失更规范的做法是继承tf.keras.losses.Loss类class CustomLoss(tf.keras.losses.Loss): def __init__(self, namecustom_loss, **kwargs): super().__init__(namename, **kwargs) def call(self, y_true, y_pred): return tf.reduce_mean(tf.square(y_true - y_pred))类写法的好处是可以带超参数、可以序列化、可以自动处理sample_weight和reduction。如果你的损失需要配置参数比如 focal loss 的 gamma强烈建议用类写法。函数式写法没法保存超参数模型保存再加载后损失就“失忆”了。2.4 三种写法的对比写法适用场景优点缺点内置损失常规任务稳定、功能全无法定制函数式简单自定义写起来快不能带参数、难序列化类继承复杂自定义可配置、可序列化代码稍多选哪种我的经验是只要损失带参数或者你打算保存模型后重新加载就用类继承。函数式只适合临时实验。3. 多损失组合的三种主流方案好进入正题。Keras 里实现多损失我总结下来有三种主流方案各有适用场景。3.1 方案一输出层命名 loss 字典这是 Keras 官方最推荐的方式。前提是你的模型有多个输出层每个输出层有自己的名字。编译时传一个字典把输出层名字映射到对应的损失函数model tf.keras.Model(inputsinputs, outputs{ classification: cls_output, regression: reg_output }) model.compile( optimizeradam, loss{ classification: categorical_crossentropy, regression: mse }, loss_weights{ classification: 1.0, regression: 0.5 } )这种方式的优雅之处在于Keras 会自动把每个输出层的预测结果和对应的标签匹配你不需要手动拆分。loss_weights参数还能控制每个损失的权重非常方便。但有个限制这种方式要求模型必须是多输出结构。如果你的模型只有一个输出但你想在同一个输出上叠加多个损失比如一个预测值同时算 MSE 和一个自定义的平滑项这个方案就不适用了。3.2 方案二自定义总损失函数如果模型只有一个输出或者你想把多个损失揉在一起可以写一个总损失函数内部调用多个子损失class CombinedLoss(tf.keras.losses.Loss): def __init__(self, alpha1.0, beta0.5, **kwargs): super().__init__(**kwargs) self.alpha alpha self.beta beta self.mse tf.keras.losses.MeanSquaredError() self.mae tf.keras.losses.MeanAbsoluteError() def call(self, y_true, y_pred): loss_mse self.mse(y_true, y_pred) loss_mae self.mae(y_true, y_pred) return self.alpha * loss_mse self.beta * loss_mae这种写法的好处是对外只暴露一个损失编译时直接传这个类就行训练代码完全不用改。缺点是你没法单独监控每个子损失的收敛情况除非在call里加tf.print或者用add_metric。我个人的偏好是如果子损失之间权重固定、不需要单独监控用这个方案最省事。如果需要动态调权重或者要单独看每个损失就得用方案三。3.3 方案三多输出 自定义损失 指标分离这是最灵活的方案结合了前两者的优点。模型多输出每个输出配一个自定义损失同时通过metrics参数单独监控每个子损失model.compile( optimizeradam, loss{ main: MainLoss(), aux: AuxLoss() }, loss_weights{main: 1.0, aux: 0.3}, metrics{ main: [accuracy], aux: [mae] } )这种方案适合复杂的多任务学习场景。每个任务的损失可以独立定制权重可以独立调整训练过程中还能看到每个任务的指标变化。代价是模型结构会复杂一些需要设计好输出层的拆分。3.4 三种方案的选型建议方案模型结构要求灵活性代码复杂度推荐场景loss 字典多输出中低标准多任务自定义总损失单输出即可低低损失组合固定多输出自定义多输出高中复杂多任务选哪个我的判断逻辑是先看模型是不是天然多输出。如果是优先用方案一或方案三。如果模型只有一个输出但需要多损失用方案二。如果既要多输出又要高度定制方案三。4. 手把手实现一个多损失训练流程光说理论没意思咱们来搭一个完整的例子。假设要做一个图像任务模型同时输出分类结果和回归结果比如预测图像类别和某个连续属性分类用交叉熵回归用自定义的平滑 L1 损失再加一个辅助的正则损失。4.1 定义各个自定义损失先写三个损失类class SmoothL1Loss(tf.keras.losses.Loss): def __init__(self, delta1.0, namesmooth_l1, **kwargs): super().__init__(namename, **kwargs) self.delta delta def call(self, y_true, y_pred): diff tf.abs(y_true - y_pred) small 0.5 * tf.square(diff) large self.delta * (diff - 0.5 * self.delta) return tf.reduce_mean(tf.where(diff self.delta, small, large)) class AuxRegularizationLoss(tf.keras.losses.Loss): def __init__(self, weight0.01, nameaux_reg, **kwargs): super().__init__(namename, **kwargs) self.weight weight def call(self, y_true, y_pred): return self.weight * tf.reduce_mean(tf.square(y_pred))分类损失直接用内置的CategoricalCrossentropy不用自己写。这里解释一下 SmoothL1 的设计意图当预测误差小于 delta 时用平方损失梯度平滑误差大于 delta 时用线性损失对异常值更鲁棒。这是从 Faster R-CNN 里借鉴过来的思路在回归任务里很常用。delta 取 1.0 是经验值你也可以根据标签的尺度调整。4.2 搭建多输出模型用函数式 API 搭一个简单的多输出模型def build_model(input_shape, num_classes): inputs tf.keras.Input(shapeinput_shape, nameinput) x tf.keras.layers.Conv2D(32, 3, activationrelu)(inputs) x tf.keras.layers.MaxPooling2D()(x) x tf.keras.layers.Conv2D(64, 3, activationrelu)(x) x tf.keras.layers.GlobalAveragePooling2D()(x) shared tf.keras.layers.Dense(128, activationrelu, nameshared)(x) cls_output tf.keras.layers.Dense(num_classes, activationsoftmax, nameclassification)(shared) reg_output tf.keras.layers.Dense(1, nameregression)(shared) return tf.keras.Model(inputsinputs, outputs[cls_output, reg_output])关键点是每个输出层都要有name这个名字后面编译时要用到。名字起得清晰一点别用output1、output2这种过两天自己都忘了哪个是哪个。4.3 编译时配置多损失model build_model(input_shape(64, 64, 3), num_classes10) model.compile( optimizertf.keras.optimizers.Adam(learning_rate1e-3), loss{ classification: tf.keras.losses.CategoricalCrossentropy(), regression: SmoothL1Loss(delta1.0) }, loss_weights{ classification: 1.0, regression: 0.5 }, metrics{ classification: [accuracy], regression: [mae] } )loss_weights这里给回归损失 0.5 的权重是因为回归损失的数值量级通常比分类损失大不降权的话回归损失会主导梯度。这个权重不是拍脑袋定的我一般会先跑几百步看两个损失的数量级然后调整到它们对总损失的贡献大致相当。如果分类损失在 2.0 左右回归损失在 10.0 左右那回归权重给 0.2 左右比较合适。4.4 准备数据管道多输出模型的数据生成器要返回一个字典键是输出层名字值是对应的标签def data_generator(batch_size): while True: images np.random.rand(batch_size, 64, 64, 3).astype(np.float32) cls_labels np.random.randint(0, 10, size(batch_size,)) cls_labels tf.keras.utils.to_categorical(cls_labels, 10) reg_labels np.random.rand(batch_size, 1).astype(np.float32) yield images, { classification: cls_labels, regression: reg_labels }这里最容易踩的坑是字典的键必须和输出层的名字完全一致大小写都不能错。我见过有人写成Classification结果 Keras 报错说找不到对应的输出排查半天。4.5 启动训练并观察各损失train_ds tf.data.Dataset.from_generator( lambda: data_generator(32), output_signature( tf.TensorSpec(shape(None, 64, 64, 3), dtypetf.float32), { classification: tf.TensorSpec(shape(None, 10), dtypetf.float32), regression: tf.TensorSpec(shape(None, 1), dtypetf.float32) } ) ) history model.fit(train_ds, steps_per_epoch100, epochs10)训练日志里会分别显示classification_loss、regression_loss和total_loss还有各自的指标。这样你就能清楚地看到每个任务的收敛情况哪个任务拖后腿一目了然。提示如果某个任务的损失一直不降先检查它的权重是不是太小再检查标签和预测的形状是否匹配。形状不匹配时 Keras 有时不会报错而是静默广播导致损失计算错误。5. 进阶技巧动态权重与损失平衡多损失训练最头疼的问题就是权重怎么定。固定权重靠调参费时费力。下面聊几个进阶技巧。5.1 用回调函数动态调整损失权重Keras 的Callback机制可以在训练过程中动态修改loss_weights。思路是监控每个损失的下降速度下降慢的任务加大权重class DynamicLossWeight(tf.keras.callbacks.Callback): def __init__(self, patience3, factor1.1): super().__init__() self.patience patience self.factor factor self.wait 0 self.best {} def on_epoch_end(self, epoch, logsNone): logs logs or {} for key in [classification_loss, regression_loss]: current logs.get(key) if current is None: continue if key not in self.best or current self.best[key]: self.best[key] current self.wait 0 else: self.wait 1 if self.wait self.patience: # 该损失停滞了加大它的权重 name key.replace(_loss, ) self.model.loss_weights[name] * self.factor self.wait 0这个思路来自“不确定性加权”的简化版。核心逻辑是哪个任务学得慢就给它更多关注。实测下来在任务难度差异大的场景里比固定权重收敛更快。5.2 基于不确定性的自动权重更学术的做法是用同方差不确定性来加权。每个任务学一个可训练的参数log_var权重是1 / (2 * exp(log_var))同时加一个log_var的正则项class UncertaintyWeightedLoss(tf.keras.layers.Layer): def __init__(self, num_tasks, **kwargs): super().__init__(**kwargs) self.num_tasks num_tasks def build(self, input_shape): self.log_vars self.add_weight( namelog_vars, shape(self.num_tasks,), initializerzeros, trainableTrue ) def call(self, losses): total 0.0 for i, loss in enumerate(losses): precision tf.exp(-self.log_vars[i]) total precision * loss self.log_vars[i] return total这个方法的妙处在于权重是学出来的不用手动调。缺点是增加了可训练参数训练初期可能不稳定。我一般会先用固定权重跑通再换成这个方案做精调。5.3 梯度归一化思路还有一种思路是 GradNorm通过调整权重让各任务的梯度范数保持相近。实现起来比较复杂需要自定义训练循环。如果你的任务差异特别大固定权重怎么调都不行可以考虑这个方向。但对大多数项目来说前两种方案够用了。5.4 权重策略对比策略实现难度调参成本适用场景固定权重低高任务量级相近动态回调中低任务难度差异大不确定性加权中低追求自动化GradNorm高中极端不平衡我的建议是先用固定权重跑基线观察各损失的数量级和收敛速度再决定要不要上动态方案。别一上来就搞最复杂的容易把自己绕进去。6. 踩坑实录与排查清单这部分是我在实际项目里踩过的坑整理成清单希望能帮你少走弯路。6.1 常见报错与解决报错信息原因解决方法Unknown loss function加载模型时找不到自定义损失用custom_objects传进去No gradients provided损失里用了不可导操作检查是否用了argmax、round等Shapes must be equal标签和预测形状不匹配检查输出层维度和标签维度KeyError: xxx字典键和输出层名不一致核对输出层name参数损失为 NaN数值溢出或除零加epsilon、用from_logitsTrue6.2 模型保存与加载的坑自定义损失在保存模型时有个大坑model.save()默认不保存自定义损失函数的代码只保存名字。加载时如果不提供custom_objects就会报Unknown loss functionmodel tf.keras.models.load_model( model.h5, custom_objects{ SmoothL1Loss: SmoothL1Loss, AuxRegularizationLoss: AuxRegularizationLoss } )更稳妥的做法是用SavedModel格式或者把损失函数注册到 Keras 的全局命名空间。我现在的习惯是所有自定义损失都放在一个独立的模块里加载时统一 import避免遗漏。6.3 损失权重调参的实操心得调权重这件事我的经验是分三步走第一步先让所有权重为 1.0跑 100 到 200 步记录每个损失稳定后的数值。这一步的目的是摸清各损失的数量级。第二步按数量级反比调整权重。比如分类损失稳定在 2.0回归损失稳定在 20.0那回归权重给 0.1让两者对总损失的贡献接近。第三步微调。观察哪个任务的指标提升慢适当加大它的权重但每次调整幅度别超过 2 倍否则容易震荡。注意不要频繁调整权重。每次调整后至少跑一个完整的 epoch 再判断效果否则你看到的是噪声不是信号。6.4 多损失训练的监控技巧训练多损失模型时光看总损失是不够的必须单独看每个子损失。我的做法是在metrics里把每个子损失也加进去metrics{ classification: [ accuracy, tf.keras.metrics.CategoricalCrossentropy(namecls_ce) ], regression: [ mae, tf.keras.metrics.MeanSquaredError(namereg_mse) ] }这样日志里能同时看到损失和指标的变化趋势。如果某个任务的损失在降但指标不升可能是过拟合或者指标和损失不匹配如果损失和指标都不动那大概率是梯度没传过去或者权重太小。6.5 一个容易被忽略的细节损失函数的 reductionKeras 的Loss基类有个reduction参数默认是SUM_OVER_BATCH_SIZE。如果你在call里已经做了reduce_mean再叠加默认的 reduction 就会导致损失被重复平均。正确做法是在__init__里显式指定super().__init__(namename, reductiontf.keras.losses.Reduction.NONE, **kwargs)然后在call里自己控制 reduction。这个坑很隐蔽因为损失值看起来“正常”只是比预期小了一个 batch size 的倍数不仔细看发现不了。7. 一些实战中的个人体会聊了这么多技术细节最后说几点我在实际项目里攒下的体会。关于代码组织我现在习惯把所有的自定义损失放在一个losses.py文件里每个损失一个类文件顶部统一 import TensorFlow。这样模型定义、训练脚本、损失函数三者解耦改哪个都不影响其他。以前我把损失写在训练脚本里后来项目大了同一个损失要在多个脚本里用复制粘贴改得到处都是维护起来很痛苦。关于调试多损失模型出问题时我第一件事是把所有权重设成 1.0只保留一个损失确认单任务能正常收敛。然后再逐个加回其他损失。这样能快速定位是哪个损失引入的问题。直接上来就调多损失出了问题根本不知道是谁的锅。关于权重初始化多任务模型里共享层的初始化很重要。如果共享层初始化不好某个任务可能一开始就占据主导其他任务永远学不起来。我的做法是共享层用 He 初始化各任务头用较小的初始化尺度让各任务起点接近。关于学习率多损失模型对学习率比单损失更敏感。因为总损失的曲率是各损失曲率的叠加学习率太大容易震荡。我一般会把学习率调得比单任务小一半左右配合余弦退火或者 ReduceLROnPlateau 使用。关于评估训练时看损失评估时一定要看各任务的实际指标。损失下降不代表业务指标提升尤其是自定义损失和业务指标不完全对齐的时候。我见过损失降得很漂亮但分类准确率纹丝不动的案例最后发现是损失权重把回归任务压得太狠分类任务根本没学到东西。这些经验不一定适用于所有场景但大方向应该是对的。多损失函数这件事说到底是在多个目标之间找平衡技术手段只是工具真正重要的是你对业务目标的理解——哪个任务更重要哪个任务可以妥协这些判断比调参技巧更关键。
