AI损失函数踩坑记录

别急着给AI损失函数下定义,先看这次卡在哪。

从常见的损失函数到自定义损失,从踩坑到优化,把实践中遇到的问题和解决方案整理出来。

从 MSE 开始

MSE(Mean Squared Error,均方误差)是最基础的损失函数之一,也是很多人写第一个模型时会接触到的。公式很直观:

def mse_loss(predictions, targets):
    return ((predictions - targets) ** 2).mean()

用 PyTorch 的实现更简单:

import torch
import torch.nn as nn

mse_loss = nn.MSELoss()
loss = mse_loss(predictions, targets)

但我第一次在实践中用 MSE 遇到的坑是离群值敏感问题。当时在做房价预测任务,数据里有几个异常高的价格(比如录入错误、特殊豪宅之类),结果模型被这几个点带着跑,整体预测效果反而变差了。

# 模拟一下离群值的影响
import numpy as np
import matplotlib.pyplot as plt

# 正常数据
normal_prices = np.array([300, 350, 400, 380, 420, 450, 410, 390]) * 10000
# 加上离群值
prices_with_outlier = np.append(normal_prices, [800 * 10000])

# MSE 对离群值的惩罚是平方级的
mse_normal = np.mean(normal_prices ** 2)
mse_outlier = np.mean(prices_with_outlier ** 2)

print(f"正常数据的 MSE: {mse_normal:.2e}")
print(f"含离群值的 MSE: {mse_outlier:.2e}")
print(f"MSE 放大倍数: {mse_outlier / mse_normal:.2f}")

输出结果大概是这样(具体数值会变,但趋势很明显):

正常数据的 MSE: 1.44e+11
含离群值的 MSE: 4.17e+11
MSE 放大倍数: 2.89

几个离群值就把 MSE 放大近 3 倍,模型自然会被带着跑。

解决方式有几种:

  1. 数据清洗,先把离群值处理掉(预处理的活)
  2. 换损失函数,比如用 MAE(Mean Absolute Error)
  3. 用 Huber Loss 这类混合损失

MAE 对离群值就没那么敏感:

def mae_loss(predictions, targets):
    return torch.abs(predictions - targets).mean()

实测房价预测任务里,把 MSE 换成 MAE 后,模型的泛化能力明显提升了,不过收敛速度慢了一点。这是 trade-off。

分类任务的损失函数

做分类任务时,Cross Entropy(交叉熵)是主流选择。PyTorch 里可以直接用 nn.CrossEntropyLoss,它内置了 LogSoftmax + NLLLoss 的组合。

import torch
import torch.nn as nn

# 分类损失
ce_loss = nn.CrossEntropyLoss()

# 注意:nn.CrossEntropyLoss 期望的输入是没有 softmax 的 logits
logits = torch.randn(8, 10)  # batch_size=8, num_classes=10
targets = torch.randint(0, 10, (8,))

loss = ce_loss(logits, targets)

这里有个坑容易踩:nn.CrossEntropyLoss 的输入应该是 logits(没有经过 softmax),而不是 softmax 后的输出。我第一次用的时候不小心传了 softmax 后的结果,模型训练了几个小时都在原地打转。

如果数据类别不平衡(比如正负样本 1:9),直接用 CrossEntropy 会偏向多数类。有几种应对方式:

  1. 加 class weights:
# 假设正样本数量少,给更高的权重
class_weights = torch.tensor([1.0, 9.0])
weighted_ce_loss = nn.CrossEntropyLoss(weight=class_weights)
  1. 用 Focal Loss:
class FocalLoss(nn.Module):
    def __init__(self, alpha=1, gamma=2, reduction='mean'):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction

    def forward(self, inputs, targets):
        ce_loss = nn.functional.cross_entropy(inputs, targets, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss

        if self.reduction == 'mean':
            return focal_loss.mean()
        elif self.reduction == 'sum':
            return focal_loss.sum()
        else:
            return focal_loss

Focal Loss 的核心思想是:让模型更关注难分类的样本(通过 (1-pt)^gamma 调整权重)。gamma 越大,难样本的权重就越高。

我在一个医疗影像分类任务里用 Focal Loss,正负样本比约 1:15,用普通 CrossEntropy 模型基本学不到正样本的特征。换成 Focal Loss(alpha=0.25, gamma=2)后,正样本的 recall 从 0.3 提升到了 0.7。

自定义损失函数

有些场景下,内置的损失函数就是不够用。我遇到过一个场景:需要同时优化多个目标,比如既要分类准确,又要预测某个属性值。这时候就需要自定义损失。

先看一个多目标学习的基本写法:

class MultiTaskLoss(nn.Module):
    def __init__(self, task_weights=None):
        super(MultiTaskLoss, self).__init__()
        self.task_weights = task_weights or [1.0, 1.0]
        self.classification_loss = nn.CrossEntropyLoss()
        self.regression_loss = nn.MSELoss()

    def forward(self, class_logits, class_targets, reg_preds, reg_targets):
        loss_cls = self.classification_loss(class_logits, class_targets)
        loss_reg = self.regression_loss(reg_preds, reg_targets)

        total_loss = (
            self.task_weights[0] * loss_cls +
            self.task_weights[1] * loss_reg
        )

        return total_loss, loss_cls, loss_reg

用起来也很简单:

multi_loss = MultiTaskLoss(task_weights=[1.0, 0.5])

total_loss, loss_cls, loss_reg = multi_loss(
    class_logits, class_targets,
    reg_preds, reg_targets
)

total_loss.backward()

但这里有个问题:任务权重怎么定?一开始我设的是 [1.0, 1.0],结果发现分类损失下降很快,但回归损失几乎不动,原因是两个 loss 的量级不一样。调整权重后效果好了不少,但这属于"调参"的范畴,没有通用解。

另一个更实用的自定义损失案例是:在 NER(命名实体识别)任务里,给错误的边界预测加惩罚。比如把一个完整的实体拆成两半,或者把两个实体合并成一个,这种错误比单纯漏掉某个实体更严重。

class NERLoss(nn.Module):
    def __init__(self, label_smoothing=0.1, boundary_penalty=2.0):
        super(NERLoss, self).__init__()
        self.ce_loss = nn.CrossEntropyLoss(label_smoothing=label_smoothing, ignore_index=-100)
        self.boundary_penalty = boundary_penalty

    def forward(self, logits, targets):
        # 基础交叉熵损失
        base_loss = self.ce_loss(logits, targets)

        # 检测边界错误并加惩罚
        # 这里简化处理,实际实现需要更复杂的逻辑
        boundary_loss = 0
        predictions = logits.argmax(dim=-1)

        # ... 检测边界错误的逻辑 ...
        # 如果检测到严重的边界错误,就加额外惩罚

        return base_loss + self.boundary_penalty * boundary_loss

这个边界惩罚不是直接写在 loss 公式里,而是在计算 loss 的过程中识别出错误模式后额外加权重。实际效果上,模型更倾向于产生"要么对要么完全不对"的预测,减少了边界模糊的情况。

损失函数的选择策略

选损失函数的时候,我通常会考虑这几个因素:

  1. 任务类型:回归、分类、多标签、生成任务等。类型不同,基础选择也不同。

  2. 数据特性:类别是否平衡、是否有离群值、标签噪声情况等。类别不平衡就要考虑加权重或 Focal Loss,离群值多就考虑 MAE 或 Huber Loss。

  3. 业务目标:有时业务要求更关注某些特定错误类型。比如医疗诊断里,漏诊的代价可能比误诊更高,这种情况下就需要在 loss 里加针对性调整。

  4. 模型特性:有些模型本身对 loss 的选择就比较敏感。比如 GAN 里,生成器和判别器的 loss 需要平衡,调不好的话就容易崩溃。

这里有个简单的决策流程,虽然不能覆盖所有情况,但能提供一些参考:

graph TD A[开始选择损失函数] --> B{任务类型} B -->|回归| C{是否有离群值} B -->|分类| D{类别是否平衡} B -->|多任务| E[考虑多目标加权] B -->|生成任务| F{是否 GAN} C -->|有| G[使用 Huber Loss 或 MAE] C -->|无| H[使用 MSE] D -->|平衡| I[使用 Cross Entropy] D -->|不平衡| J[考虑 Focal Loss 或加 class weights] E --> K[多目标加权 + 自适应权重] F -->|是| L[Wasserstein Loss + Gradient Penalty] F -->|否| M[根据具体生成任务选择] G --> N[验证效果] H --> N I --> N J --> N K --> N L --> N M --> N

优化技巧和坑点

1. Loss 数值的监控

训练过程中,只看 total loss 往往不够。我习惯把每个子任务的 loss 单独打出来,这样能看出哪个任务出了问题。

def train_epoch(model, dataloader, optimizer, loss_fn):
    model.train()
    total_loss = 0
    task_losses = []  # 存储每个 task 的 loss

    for batch in dataloader:
        optimizer.zero_grad()

        # 假设 loss_fn 返回 (total_loss, loss_cls, loss_reg)
        loss, loss_cls, loss_reg = loss_fn(model(batch))

        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        task_losses.append([loss_cls.item(), loss_reg.item()])

    avg_total = total_loss / len(dataloader)
    avg_task_losses = np.mean(task_losses, axis=0)

    return avg_total, avg_task_losses

这样就能看出哪个 task 的 loss 没降下来,或者降得太快导致其他 task 学不到东西。

2. Loss 的数值范围

有些损失函数的数值范围差得很多,直接加权容易出问题。比如 CrossEntropyLoss 的输出大概在 0-5 之间,但 MSE 的输出可能是几千上万。这种情况下,直接给两个 loss 各配 0.5 的权重,实际效果是 MSE 占主导。

解决方式有几种:

  1. 归一化:把每个 loss 除以它的初始值或移动平均值
  2. 自适应权重:用某种机制自动调整权重(比如 GradNorm、Uncertainty Weighting)
  3. 经验调整:先定个初始权重,根据训练效果微调

自适应权重的一个简单实现(Uncertainty Weighting):

class UncertaintyWeightedMultiTaskLoss(nn.Module):
    def __init__(self, num_tasks):
        super().__init__()
        # 学习任务权重(用 log 参数化,保证为正)
        self.log_vars = nn.Parameter(torch.zeros(num_tasks))

    def forward(self, *losses):
        losses = torch.stack(losses)
        weighted_losses = losses / (2 * self.log_vars.exp())
        total_loss = weighted_losses.sum() + self.log_vars.sum()
        return total_loss

3. 梯度问题

有些损失函数容易出现梯度消失或爆炸问题。比如用 MSE 做 cross entropy 的近似时,预测值和真实值差距太大时,梯度可能会爆炸。

# 容易出问题的写法
def bad_cross_entropy_approx(logits, targets):
    probs = torch.softmax(logits, dim=-1)
    return torch.mean((probs - targets) ** 2)  # 用 MSE 近似 cross entropy

梯度爆炸的处理方式:

  1. 梯度裁剪:nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
  2. 调整学习率
  3. 换损失函数,用更稳定的实现

4. 验证集上的 loss 验证

有时候训练集 loss 降得很漂亮,但验证集 loss 反而上升。这通常是过拟合,但也可能是验证集和训练集的分布差异太大。

我遇到过一次情况:训练集 loss 持续下降,验证集 loss 在某个点后突然跳起来。排查后发现是验证集里有几个异常样本,它们的特征和训练集完全不一样,模型根本没见过这种数据。

解决方案:

  1. 检查数据分布,确保训练集和验证集一致
  2. 用更鲁棒的损失函数(比如 Huber Loss)
  3. 数据增强,让训练集覆盖更多情况

实际案例:目标检测的损失调整

去年做一个工业缺陷检测项目,需要在流水线上实时检测产品瑕疵。一开始用的是标准的 YOLOv5,直接用它自带的损失函数。

跑了几轮后发现一个问题:模型对小目标的检测效果很差。一个小缺陷可能只占图像的 1% 不到,标准的 IoU 损失对小目标的变化不太敏感。

# 标准 IoU loss 的问题
def standard_iou_loss(pred_boxes, gt_boxes):
    iou = calculate_iou(pred_boxes, gt_boxes)
    return 1 - iou.mean()

# 小目标场景下,IoU 的变化对 loss 的影响很小
# 比如:小目标 IoU 从 0.8 变到 0.6,loss 变化只有 0.2
# 但这对检测效果影响很大

我们的解决方案是:

  1. 给小目标更高的权重
  2. 加 Focal Loss 的思想,让模型更关注难检测的样本
class WeightedIoULoss(nn.Module):
    def __init__(self, small_object_threshold=0.01, small_object_weight=2.0):
        super().__init__()
        self.small_object_threshold = small_object_threshold
        self.small_object_weight = small_object_weight

    def forward(self, pred_boxes, gt_boxes):
        iou = calculate_iou(pred_boxes, gt_boxes)
        base_loss = 1 - iou

        # 计算目标面积
        box_areas = calculate_box_area(gt_boxes)
        normalized_areas = box_areas / box_areas.max()

        # 小目标加权
        weights = torch.where(
            normalized_areas < self.small_object_threshold,
            self.small_object_weight,
            1.0
        )

        return (weights * base_loss).mean()

调整后,小目标的检测 mAP 从 0.45 提升到了 0.68,虽然大目标的 mAP 稍微降了一点(从 0.92 到 0.89),但整体效果好了很多。

收尾与思考

损失函数这东西,说简单也简单,说复杂也复杂。简单在于大多数任务都有现成可用的选择;复杂在于当现成选择不够用时,如何根据具体场景调整。

我踩过的坑主要集中在这几个地方:

  1. 任务类型和数据特性没先搞清楚,直接用了标准损失函数
  2. 多任务学习时,权重定得太随意,导致某些任务学不到东西
  3. 没有监控子任务 loss,出问题时不知道是哪里不对
  4. 过度追求 loss 的数学完美,忽视了实际效果

有时候,一个稍微"糙"一点的损失函数,配合合理的调参,反而比精心设计的复杂 loss 更好用。工程里太多事情都是这样:够用就好。

如果你也在调损失函数的过程中遇到过奇怪的问题,不妨分享一下。毕竟,踩坑也是技术积累的一部分。

版权声明: 本文首发于 指尖魔法屋-AI损失函数踩坑记录https://blog.thinkmoon.cn/post/237-loss-functions-mse-custom-optimization/) 转载或引用必须申明原指尖魔法屋来源及源地址!