在Keras里玩到第二阶段大概率会撞上一堵墙内置的损失函数不够用了。想做类别不平衡的分类内置的categorical_crossentropy对少数类毫无办法想给自编码器加一个结构正则内置loss列表里根本没有那种“MSE加KL散度”的组合项想搞对比学习或者三元组损失那更是只能自己手搓。于是“自定义loss函数”就成了绕不过去的必修课。这篇东西就围绕“怎么写一个能正常训练、能保存模型、算完梯度还不出NaN的自定义loss函数”来展开。我会先讲清楚Keras到底把loss函数当成什么来调用再给几个可以直接抄的实现然后重点说说保存、序列化、访问中间层这些进阶需求对应的坑最后用一次实际排查经历收尾。适合已经会跑通基本模型的同学也适合被奇怪bug卡住半天的朋友。1. 自定义loss的本质Keras到底怎么调用你的损失函数1.1 签名只有一个y_true和y_pred先看最朴素的写法import tensorflow as tf def my_loss(y_true, y_pred): return tf.reduce_mean(tf.square(y_true - y_pred))这几乎就是MSE。Keras训练时是这么对待它的拿一批数据过一遍模型得到y_predy_true是从数据里取出来的标签然后调用my_loss(y_true, y_pred)得到一个标量或者一个形状为(batch,)的张量再对这个张量求平均得到最终的标量loss最后对loss求梯度反向传播更新参数。所以自定义loss本质上不是“你发明了一个新公式”而是“你给Keras提供了一台从y_true和y_pred到损失张量的函数机器”。Keras要求这个函数的输入输出满足固定约定第一个参数是真实值第二个是预测值函数体内部必须全部是可微的张量运算输出只要是个能求均值的东西即可。这里有个不少人会忽略的点Keras并不会强制你返回标量。你可以返回一个形状为(batch,)的loss向量框架会在外面帮你做tf.reduce_mean。甚至在某些场景里比如你想给每一条样本不同的权重你更希望返回向量然后自己在函数内部做加权平均。这跟内置损失类的reduction参数有对应关系但自定义普通函数时没有直接的reduction参数你得自己在函数内部决定怎么合并。用生活类比来说框架像是一个包工头他规定每个工人loss函数必须从A手上领到y_true、从B手上领到y_pred最后交回一个能算损失的东西。至于你用加法、乘法、除法还是更奇怪的组合他不管。这既是自由也是责任——你交回去的东西如果包含NaN、inf或者形状和y_true根本对不上他也不会提前提示。1.2 三个容易被忽略的暗坑第一个暗坑是自定义loss里不能有Python的if/else去决定不同的计算路径也不能用numpy的数组操作。你在函数里写的所有逻辑最终都要被追踪求导所以必须用TensorFlow算子。想按条件分支请用tf.where想取索引请用tf.gather想求最大值请用tf.reduce_max而不是np.max。一旦两者混着写训练时要么报错要么梯度路径断裂loss看起来在降但模型权重根本没怎么更新。第二个暗坑是y_true的形状未必是你以为的形状。分类任务里如果loss直接吃softmax输出y_true是one-hot向量形状是(batch, num_classes)如果用了sparse_categorical_crossentropyy_true是整数索引形状是(batch,)——你在自定义loss里就要用tf.one_hot先扩展。遇到这种情况最好在写函数之前先用一个固定形状的假张量手动调一遍打印一下shape再安心接入compile。第三个暗坑是mask。给RNN或者Embedding设置mask_zeroTrue时Keras内置loss会自动忽略被padding的0位置但换成自定义loss之后自动mask并不总是生效。Keras 3对loss的mask处理做了不少改动可只要你在自定义函数里用了tf.reduce_mean它就会把padding位置也算进去导致训练指标明显变差。一个稳妥的做法是在自定义loss里拿到序列的真实长度或者调用tf.boolean_mask提前把无效位置过滤掉。2. 三个可直接抄的loss实现从加权交叉熵到组合式正则损失2.1 类别不平衡内置class_weight不够时加权交叉熵怎么写model.fit(..., class_weight...)确实可以给不同类别加权但它在语义上只是样本权重而且对自定义训练循环、多输出模型作用域非常有限。更可控的做法是把权重写进loss闭包里import tensorflow as tf from tensorflow import keras def weighted_categorical_crossentropy(weights): # 假设类别数len(weights)比如[1.0, 5.0, 2.0] weights tf.constant(weights, dtypetf.float32) def loss(y_true, y_pred): y_pred tf.clip_by_value(y_pred, keras.backend.epsilon(), 1.0 - keras.backend.epsilon()) # y_true是one-hot时 per_sample -tf.reduce_sum(weights[None, :] * y_true * tf.math.log(y_pred), axis-1) return tf.reduce_mean(per_sample) return loss model.compile(optimizeradam, lossweighted_categorical_crossentropy([1.0, 5.0, 2.0]))这里用到了闭包技巧外层函数接收一个普通Python列表把数值冻结成常量内层函数才是Keras真正调用的loss(y_true, y_pred)。权重是固定常量训练过程中不会被梯度影响。有个细节需要注意weights[None, :]是把一维权重扩展成形状(1, num_classes)目的是让它跟y_true的(batch, num_classes)广播。类权重的值不要随手拍建议用多数类样本数/少数类样本数的比值或者反过来看你的优化目标。实测里权重从1倍加到5倍少数类的召回率会显著提升但精确率会掉如果已经用了focal loss这类自适应加权就不必再叠一个很重的固定权重。2.2 Focal Lossalpha和gamma怎么落进函数类别不平衡严重、且存在大量易分样本时比如检测任务里的背景框、分割图里大面积背景固定加权交叉熵并不够用。大家通常会转向Focal Loss思路很简单让损失函数自动聚焦到难分样本上——难分样本梯度大易分样本梯度小。def focal_loss(alpha0.25, gamma2.0): def loss(y_true, y_pred): y_pred tf.clip_by_value(y_pred, keras.backend.epsilon(), 1.0 - keras.backend.epsilon()) y_true tf.cast(y_true, tf.float32) # 标准交叉熵 ce -y_true * tf.math.log(y_pred) - (1 - y_true) * tf.math.log(1 - y_pred) # 调制因子p_t越接近1越是易分样本因子越小 p_t y_true * y_pred (1 - y_true) * (1 - y_pred) modulating_factor tf.pow(1.0 - p_t, gamma) # alpha是对正负样本的兜底平衡 alpha_t y_true * alpha (1 - y_true) * (1 - alpha) return tf.reduce_mean(alpha_t * modulating_factor * ce) return loss注意这个实现假定y_true是0/1的二分类标签输出层用Dense(1, activationsigmoid)。多分类版本需要把y_true换成one-hot、把1-y_true改成逐类构造公式会复杂不少。用二分类版接在sigmoid后面是最省心的状态。gamma取2.0是论文里的默认值但不代表所有场景都最优。我自己调参时发现如果模型过拟合很快、训练loss几乎不降把gamma降到0.5到1.0反而更稳定反过来如果正负样本极端不平衡gamma可以调到2.5。alpha的作用是“兜底”——当调制因子把所有易分样本都压得极低时少数类样本会变得过于微弱alpha正好起对冲作用。2.3 自编码器里的“重构正则”组合loss自编码器这类模型常常需要组合多个目标重构误差要小、隐空间分布要符合先验。你可以尝试写一个组合函数def ae_loss(alpha1.0, beta0.1): def loss(y_true, y_pred): # 这里能拿到的只有重建的pred拿不到encoder的隐变量 reconstruction tf.reduce_mean(tf.square(y_true - y_pred)) kl 0.0 # 想加KL项但这里根本没有z_mean、z_log_var return alpha * reconstruction beta * kl return loss这里就会暴露纯自定义loss的局限它跟模型内部结构是隔离的只有y_true和y_pred两个参数。要让loss用上中间层的变量你需要换一个思路——用model.add_loss这个我在第4节会细讲。“重构正则”在Keras里更正统的写法是模型主体是input - encoder - decoder - output同时把编码器输出的z_mean、z_log_var作为模型额外输出再用add_loss或者自定义Metric来处理KL项。很多人一上来就写一个奇怪的自定义loss把z_mean塞在y_pred的某个下标里这种方案能用但真的不建议因为它会把模型输出定义搅得很乱推理时还得从乱七八糟的分支里取数据。3. loss必须能序列化保存模型时最常见的连环坑3.1 一个可保存、可加载的自定义loss需要先注册自定义loss如果只是随手写在脚本里model.save(xxx.h5)之后再keras.models.load_model(xxx.h5)十有八九会报错报错关键词大概是Unknown loss function: xxx。原因是H5格式只存了一个函数名它需要你在加载时用custom_objects把名字映射回函数或者直接让Keras记住这个函数。推荐做法是使用Keras自带的注册装饰器import keras keras.saving.register_keras_serializable() def focal_loss(y_true, y_pred, alpha0.25, gamma2.0): ...用装饰器之后model.save和load_model就能自动找到这个函数不需要每次手动传custom_objects。这对团队协作特别友好同事拿到h5文件不用翻聊天记录找“到底哪个loss定义跑得通”。如果你用的是tf.keras对应写法是tf.keras.saving.register_keras_serializable()机制一样。但要注意闭包工厂配合序列化时比较麻烦。如果你写的是loss weighted_categorical_crossentropy([1.0, 5.0, 2.0])Keras保存下来的名字和参数绑定关系在不同版本里表现不一致旧版本甚至容易丢参数。为了少踩坑我更推荐把参数全部放进函数签名里用默认值维护keras.saving.register_keras_serializable() def weighted_categorical_crossentropy(y_true, y_pred, weightsNone): if weights is None: weights [1.0, 1.0, 1.0] weights tf.constant(weights, dtypetf.float32) ...然后compile(lossweighted_categorical_crossentropy)。如果你确实需要在不同任务上用不同权重可以在调用时包一层局部函数但序列化字段会把weights参数焊死在配置里这样加载时才能还原。3.2 compile里传字符串和传函数对象到底有什么区别model.compile(lossmse)和model.compile(losstf.keras.losses.MeanSquaredError())在多数情况下等价但这背后的映射机制会骗到人。字符串会被Keras的注册表解析成对应的类实例或函数如果你自己造了一个函数名叫focal_loss不注册的话lossfocal_loss必然报错。这是自定义loss接触到的第一个序列化暗坑。另一个很容易误会的点是内置损失类的reduction参数。MeanSquaredError(reductionsum_over_batch_size)返回的是平均后的标量reductionnone返回的是每个样本的向量。普通自定义函数没有这个参数所以Keras默认会对返回结果再过一遍均值。如果你在自定义函数内部已经reduce_mean了一次最后框架又mean一次——数值变小但梯度方向不变训练结果一般问题不大但loss曲线跟论文里对不上自己看的时候会慌。我还遇到过一种情况自定义loss返回了形状(batch, time)的序列损失忘了在时间维度上合并导致后来的tf.reduce_mean把所有时间步也平均进去。这在序列模型上的表现就是把padding位置的0误差也平均进去指标偏低但模型不崩非常隐蔽。4. 不止算y_true和y_pred如何在loss内部访问中间层特征4.1 闭包捕获、模型多输出、add_loss三选一如果你发现自定义loss只靠y_true和y_pred搞不定第一个直觉就是把中间层张量“塞”进参数里。常见的土办法是让模型输出同时包含最终logits和中间embedding然后loss里手动切分。这个办法可行但会让预测接口非常难用还得设置多个输出。更好的方案有三个闭包捕获在函数式模型里中间层张量在构建模型时已经是现成的自定义loss外面再套一层函数把中间层张量作为自由变量捕获进去。模型输出依然干净。model.add_loss直接把某个计算图节点的损失加到模型总loss里。它不需要y_true最适合正则项。子类化模型在call里计算出embedding然后显式调用损失函数并累加。闭包捕获的代码很直观def build_model(): inputs keras.Input(shape(32,)) x layers.Dense(16)(inputs) z layers.Dense(8, nameembedding)(x) outputs layers.Dense(1, activationsigmoid)(z) model keras.Model(inputs, outputs) def combined_loss(y_true, y_pred): ce keras.losses.binary_crossentropy(y_true, y_pred) kl tf.reduce_mean(tf.square(z)) # z在闭包里 return ce 0.01 * kl model.compile(optimizeradam, losscombined_loss) return model用这个方式embedding层没有作为模型输出暴露出来推理顶点依然是sigmoid输出。Keras在追踪梯度时能顺着combined_loss中的z反向传播到embedding层效果等同把z塞进y_pred。4.2 一个三元组/对比损失实例对比学习场景里自定义loss经常要同时用多个样本的embedding。最典型的是Triplet Loss需要anchor、positive、negative三个embedding。Keras的普通loss函数只有y_true和y_pred没法直接表达“我要把batch分成三组”。通常的处理方式是数据加载器把三元组拼成一个batch模型输出统一长度的embeddingloss函数内部按索引切分def triplet_loss(y_true, y_pred, margin1.0): batch_size tf.shape(y_pred)[0] # 假设按顺序每3个样本一组anchor, positive, negative anchor y_pred[0::3] positive y_pred[1::3] negative y_pred[2::3] pos_dist tf.reduce_sum(tf.square(anchor - positive), axis1) neg_dist tf.reduce_sum(tf.square(anchor - negative), axis1) loss tf.maximum(pos_dist - neg_dist margin, 0.0) return tf.reduce_mean(loss)这个实现里y_pred形状是(3N, embedding_dim)y_true在训练时根本用不到但Keras要求loss函数必须有这个参数所以照常写上内部不引用即可。关键技巧是0::3这种切片在Tensor上可以正常工作因为底层索引操作被重载了。实际使用中batch size必须是3的倍数数据迭代器必须保证每三行是一组三元组。这种强约定最好在loss里加一句断言batch_size tf.shape(y_pred)[0] tf.debugging.assert_equal(tf.math.floormod(batch_size, 3), 0, messagetriplet loss要求batch_size是3的倍数)训练初期通常没问题一旦batch size设错报错信息会直接指到loss函数能节约大量排查时间。类似思路也适用于siamese网络、知识蒸馏的师生loss等场景本质都是“需要访问比y_true和y_pred更丰富的信息”。5. 验证自定义loss的三板斧手动算一遍、数值梯度、训练冒烟5.1 用numpy先算一遍期望值自定义loss写完我几乎不会直接train。先造两组固定张量手动算一下期望输出值全程不碰模型。import numpy as np y_true_np np.array([1.0, 0.0, 1.0, 0.0]) y_pred_np np.array([0.8, 0.3, 0.6, 0.1]) loss_value focal_loss(y_truetf.constant(y_true_np), y_predtf.constant(y_pred_np)) print(loss_value.numpy())再在纸上或用numpy按公式算一遍同样的输入两边数值应该完全一致至少小数点后5位。这一步看着笨但它能一次性暴露公式写错、shape广播错、少了个负号之类的问题。不要跳过因为等模型真正训起来loss曲线复杂你根本分不清是数据问题还是loss写错。5.2 数值梯度检查验证可导性比“算对公式”更关键的是“梯度正确”。自定义loss里如果不小心用了tf.round、tf.argmax这类不可导算子公式看似能出loss但梯度为0模型永远不动。做法是用tf.GradientTape对某个假参数求梯度再和数值差分对比x tf.Variable(tf.random.normal([4, 8])) with tf.GradientTape() as tape: y_pred tf.nn.sigmoid(x) # 模拟模型输出 loss my_loss(tf.zeros_like(y_pred), y_pred) grad tape.gradient(loss, x)数值差分的关键是给每个参数维度加或减一个小ε比如1e-5算两个loss的斜率检查它和解析梯度是否接近。这一套能过基本可以放心训练。我之前写过一版对比损失里面多了个tf.argmax去对样本pair排序前向计算完全正常手动验证也对但训练时embedding根本不更新。最后靠打印梯度发现embedding层的梯度全是0顺藤摸瓜找到不可导算子。所以“能算出loss”和“能正常训练”是两回事。5.3 一个epoch的训练冒烟测试验完公式和梯度后最后再做一次冒烟测试用几行模拟数据编译模型fit一个epoch观察loss是否从预期初值开始下降。冒烟测试的数据可以故意造得很简单比如所有样本都同一个类别或者特征完全一致目的是让网络快速过拟合。如果这种简单数据上loss都不降问题大概率还是出在loss本身。冒烟测试里我还习惯把verbose1打开看第一轮loss和最后一轮loss的数量级。如果一个收敛正常的模型loss从10降到0.1而你的自定义loss跑出来是1e-5起步那大概率是scale差了一个常数模型的初始化敏感度也会随之改变必要时要调整学习率。6. 一次真实翻车记录多任务模型自定义loss的维度错配排查6.1 现象loss一路飙升然后变NaN之前我处理过一个多任务模型一个共享主干分出两个头一个做分类一个做回归。最初用的是内置loss列表model.compile( optimizeradam, loss{ cls_head: binary_crossentropy, reg_head: mse, }, loss_weights{cls_head: 1.0, reg_head: 0.1}, )训练一直正常。后来想把两个loss融合成一个自定义函数并把一个正则项加进去于是改成def multi_task_loss(y_true, y_pred): # y_pred是一个list还是拼接张量 ...结果第一个epoch loss就冲到几千然后变成NaN。当时第一反应是学习率太大、数据有问题全部排查了一遍后才意识到问题在自定义loss的输入上——多输出模型的y_pred到底是什么。6.2 排查链路拆开单分支、打印逐层shape、逐项修复多输出模型如果只传一个loss函数给compileKeras会尝试对这个函数做广播调用它会用同一个函数分别处理每个输出头的y_true和y_pred而不是把两个输出头打包成一个参数。换句话说model.compile(lossmulti_task_loss)里的multi_task_loss会被调用两次一次处理分类头一次处理回归头每次调用拿到的y_pred都只是单一头的输出张量。我当时以为它能收到所有预测的字典结果它收到的是单个张量按原逻辑算出来的loss自然只覆盖了其中一个头另一头完全失控。排查步骤可以复现如下把自定义loss改成一个打印形状的探针函数def probe_loss(y_true, y_pred): print(y_true.shape:, y_true.shape, y_pred.shape:, y_pred.shape) return tf.reduce_mean(tf.square(y_true - y_pred))fit一个epoch盯着输出。Keras会在前向阶段把每个头的shape打出来。看到分类头的y_predshape是(batch, 1)回归头也是(batch, 1)之后才明白问题在调用方式上。修复方案有两个。方案A是放弃用一个函数包两个输出改回loss{cls_head: fn1, reg_head: fn2}各写各的loss用loss_weights合并。方案B是坚持用一个自定义函数但给模型只留一个输出把分类和回归结果拼成一个(batch, 2)张量在loss内部按列切分。这是常用但会把模型输出语义搞乱的写法不太推荐。我最后选了方案A因为职责清晰、序列化也无障碍两个head各管各的loss权重叠加交给compile配置。如果非要在loss内部控制某个head参与多少权重那就把合并逻辑放在一个自定义Metric里——Metric只展示、不参与梯度思路更干净。6.3 这个坑教会我的事自定义loss里的axis永远要显式写这次翻车还带出一个隐藏问题。回归头用的是MSE正常写法def mse_loss(y_true, y_pred): return tf.reduce_mean(tf.square(y_true - y_pred))这里tf.reduce_mean会把所有维度全部平均恰好得到标量。但如果回归头输出的维度是(batch, 32, 2)比如预测32个点、每个点2个坐标而y_true也是(batch, 32, 2)这个reduce_mean会把batch、序列长度、坐标维度全部平均。从数学上没问题可一旦中间某个维度被padding填充了0或者不同时间步的误差权重应该不同你就没法精确控制。所以我现在写自定义loss的习惯是所有reduce_mean和reduce_sum必须显式写axis。例如对坐标维度求和、对时间维度平均、最后对整个batch平均。宁可多写几步也不依赖“默认把全部维度压掉”的省事操作。多任务模型尤其如此两个head的shape一旦不同就很容易被各种广播语义绕晕再让reduce隐藏维度出问题只是时间问题。把shape打印成日志把axis挑明白反而能逼着自己把loss的结构想清楚。最后说点个人体会。自定义loss是我在Keras里见过“最简单也最容易出错”的一个环节API本身极简但背后牵扯的序列化、shape广播、mask、多输出调用方式都是文档不会主动强调的知识。我的建议是每次写新loss都先走一遍第5节的三板斧numpy手算、梯度检查、单epoch冒烟。这套流程能替你挡住大部分暗坑剩下的就靠把调试信息印清楚然后耐心看shape。