【Bug已解决】PyTorch custom loss function 解决方案问题描述在深度学习中标准的损失函数如CrossEntropyLoss、MSELoss并不总是能满足所有任务的需求。许多场景需要自定义损失函数如 Focal Loss、Dice Loss、Triplet Loss、Contrastive Loss 等。然而在 PyTorch 中实现自定义损失函数时开发者经常遇到梯度无法传播、数值不稳定、维度不匹配等问题。典型的问题场景包括自定义损失函数中使用了不可导操作导致梯度为 None 或 NaN使用numpy或list操作代替 tensor 操作导致梯度断裂损失函数中数值不稳定出现inf或NaN损失函数的输入维度不匹配在损失函数中使用for循环导致性能极差自定义损失函数不支持 GPU 或多 GPU损失函数中的detach()使用不当导致梯度断裂这些问题的核心在于理解 PyTorch 的自动微分机制以及如何正确实现可导的自定义操作。错误复现场景一梯度断裂import torch import torch.nn as nn import numpy as np class WrongLoss(nn.Module): def forward(self, predictions, targets): # 错误使用 numpy 操作梯度断裂 predictions_np predictions.detach().cpu().numpy() targets_np targets.cpu().numpy() loss np.mean((predictions_np - targets_np) ** 2) return torch.tensor(loss, requires_gradTrue) model nn.Linear(10, 1) x torch.randn(5, 10) y torch.randn(5, 1) pred model(x) loss_fn WrongLoss() loss loss_fn(pred, y) loss.backward() print(model.weight.grad) # None! 梯度没有传播到模型参数场景二数值不稳定class UnstableLoss(nn.Module): def forward(self, predictions, targets): # 错误直接使用 log可能出现 log(0) -inf loss -targets * torch.log(predictions) return loss.mean() pred torch.tensor([0.0, 0.5, 1.0]) target torch.tensor([1.0, 1.0, 0.0]) loss UnstableLoss()(pred, target) # tensor([inf, -0.6931, nan])场景三维度不匹配class DimensionMismatchLoss(nn.Module): def forward(self, predictions, targets): # 错误没有处理 batch 维度 return (predictions - targets) ** 2 # 返回矩阵而不是标量 pred torch.randn(4, 10) target torch.randn(4, 10) loss DimensionMismatchLoss()(pred, target) print(loss.shape) # torch.Size([4, 10]) —— 应该是标量场景四不可导操作class NonDifferentiableLoss(nn.Module): def forward(self, predictions, targets): # 错误argmax 不可导 predicted_classes predictions.argmax(dim1) accuracy (predicted_classes targets).float().mean() # 尝试将 accuracy 作为 loss return 1.0 - accuracy # 梯度为 0 loss NonDifferentiableLoss()(pred, target) loss.backward() # 梯度全为 0因为 argmax 不可导根因分析1. PyTorch 自动微分机制PyTorch 使用动态计算图进行自动微分。只有使用torch.Tensor操作且 tensor 的requires_gradTrue时梯度才会被追踪。以下操作会断裂梯度链.detach()将 tensor 从计算图中分离.numpy()转为 numpy 数组脱离 PyTorch.item()转为 Python 标量torch.no_grad()上下文管理器中不追踪梯度不可导操作argmax、round、比较操作等2. 数值稳定性常见的数值不稳定问题log(0)-infexp(large_number)inf0 / 0NaNinf - infNaN3. 损失函数的基本要求一个正确的损失函数必须输出标量或可通过 reduction 变为标量所有操作可导或使用可导的近似数值稳定使用clamp、logsumexp等技巧支持 GPU 和多精度训练高效避免 Python 循环使用向量化操作解决方案方案一使用 tensor 操作实现自定义损失import torch import torch.nn as nn import torch.nn.functional as F class FocalLoss(nn.Module): Focal Loss - 解决类别不平衡问题 def __init__(self, alpha0.25, gamma2.0, reductionmean): super().__init__() self.alpha alpha self.gamma gamma self.reduction reduction def forward(self, predictions, targets): Args: predictions: [batch_size, num_classes] logits targets: [batch_size] class indices # 计算交叉熵 ce_loss F.cross_entropy(predictions, targets, reductionnone) # 计算 pt正确类别的概率 pt torch.exp(-ce_loss) # Focal 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 class DiceLoss(nn.Module): Dice Loss - 用于图像分割 def __init__(self, smooth1.0, reductionmean): super().__init__() self.smooth smooth self.reduction reduction def forward(self, predictions, targets): Args: predictions: [batch_size, ...] probabilities (0-1) targets: [batch_size, ...] binary labels (0 or 1) # 展平 predictions predictions.view(predictions.size(0), -1) targets targets.view(targets.size(0), -1) # 计算 Dice 系数 intersection (predictions * targets).sum(dim1) union predictions.sum(dim1) targets.sum(dim1) dice (2.0 * intersection self.smooth) / (union self.smooth) # Dice Loss 1 - Dice loss 1.0 - dice if self.reduction mean: return loss.mean() elif self.reduction sum: return loss.sum() else: return loss class TripletLoss(nn.Module): Triplet Loss - 用于度量学习 def __init__(self, margin1.0): super().__init__() self.margin margin def forward(self, anchor, positive, negative): Args: anchor, positive, negative: [batch_size, embedding_dim] # 计算距离 pos_dist F.pairwise_distance(anchor, positive, p2) neg_dist F.pairwise_distance(anchor, negative, p2) # Triplet Loss loss F.relu(pos_dist - neg_dist self.margin) return loss.mean()方案二使用数值稳定的实现class StableBCELoss(nn.Module): 数值稳定的 Binary Cross Entropy def forward(self, predictions, targets): # 使用 log-sigmoid trick 保持数值稳定 # log(sigmoid(x)) -softplus(-x) # log(1 - sigmoid(x)) -x - softplus(-x) loss -(targets * F.logsigmoid(predictions) (1 - targets) * F.logsigmoid(-predictions)) return loss.mean() class LabelSmoothingLoss(nn.Module): 标签平滑损失 def __init__(self, num_classes, smoothing0.1): super().__init__() self.num_classes num_classes self.smoothing smoothing self.confidence 1.0 - smoothing def forward(self, predictions, targets): Args: predictions: [batch_size, num_classes] logits targets: [batch_size] class indices # 创建平滑后的标签分布 with torch.no_grad(): true_dist torch.zeros_like(predictions) true_dist.fill_(self.smoothing / (self.num_classes - 1)) true_dist.scatter_(1, targets.unsqueeze(1), self.confidence) # 计算交叉熵 log_probs F.log_softmax(predictions, dim-1) loss -(true_dist * log_probs).sum(dim-1).mean() return loss方案三组合损失函数class CombinedLoss(nn.Module): 组合多种损失函数 def __init__(self, loss_weightsNone): super().__init__() self.ce_loss nn.CrossEntropyLoss() self.dice_loss DiceLoss() self.focal_loss FocalLoss() if loss_weights is None: self.weights {ce: 1.0, dice: 0.5, focal: 0.3} else: self.weights loss_weights def forward(self, predictions, targets): # 计算各种损失 ce self.ce_loss(predictions, targets) # 对于 dice loss需要概率值 probs F.softmax(predictions, dim1) # 将 targets 转为 one-hot one_hot F.one_hot(targets, num_classespredictions.size(1)).float() dice self.dice_loss(probs, one_hot) focal self.focal_loss(predictions, targets) # 加权组合 total (self.weights[ce] * ce self.weights[dice] * dice self.weights[focal] * focal) return total, {ce: ce.item(), dice: dice.item(), focal: focal.item()}完整修复代码 完整的 PyTorch 自定义损失函数实现方案 涵盖多种损失函数、数值稳定性、梯度检查、组合损失 import torch import torch.nn as nn import torch.nn.functional as F from typing import Optional, Dict, Tuple import math # # 基础损失函数 # class FocalLoss(nn.Module): Focal Loss - 解决类别不平衡 FL(p_t) -alpha * (1 - p_t)^gamma * log(p_t) def __init__(self, alpha: float 0.25, gamma: float 2.0, reduction: str mean): super().__init__() self.alpha alpha self.gamma gamma self.reduction reduction def forward(self, logits: torch.Tensor, targets: torch.Tensor) - torch.Tensor: ce_loss F.cross_entropy(logits, targets, reductionnone) 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() return focal_loss class DiceLoss(nn.Module): Dice Loss - 用于分割任务 DL 1 - (2 * |X ∩ Y| smooth) / (|X| |Y| smooth) def __init__(self, smooth: float 1.0, reduction: str mean): super().__init__()  self.smooth smooth self.reduction reduction def forward(self, predictions: torch.Tensor, targets: torch.Tensor) - torch.Tensor: predictions predictions.view(predictions.size(0), -1) targets targets.view(targets.size(0), -1) intersection (predictions * targets).sum(dim1) union predictions.sum(dim1) targets.sum(dim1) dice (2.0 * intersection self.smooth) / (union self.smooth) loss 1.0 - dice if self.reduction mean: return loss.mean() elif self.reduction sum: return loss.sum() return loss class TripletLoss(nn.Module): Triplet Loss - 度量学习 L max(d(a, p) - d(a, n) margin, 0) def __init__(self, margin: float 1.0): super().__init__() self.margin margin def forward(self, anchor, positive, negative): pos_dist F.pairwise_distance(anchor, positive, p2) neg_dist F.pairwise_distance(anchor, negative, p2) loss F.relu(pos_dist - neg_dist self.margin) return loss.mean() class ContrastiveLoss(nn.Module): 对比损失 def __init__(self, margin: float 2.0): super().__init__() self.margin margin def forward(self, embedding1, embedding2, label): label1: 相似对 label0: 不相似对 dist F.pairwise_distance(embedding1, embedding2) loss label * torch.pow(dist, 2) \ (1 - label) * torch.pow(torch.clamp(self.margin - dist, min0.0), 2) return loss.mean() class LabelSmoothingLoss(nn.Module): 标签平滑损失 def __init__(self, num_classes: int, smoothing: float 0.1): super().__init__() self.num_classes num_classes self.smoothing smoothing self.confidence 1.0 - smoothing def forward(self, logits, targets): with torch.no_grad(): true_dist torch.zeros_like(logits) true_dist.fill_(self.smoothing / (self.num_classes - 1)) true_dist.scatter_(1, targets.unsqueeze(1), self.confidence) log_probs F.log_softmax(logits, dim-1) loss -(true_dist * log_probs).sum(dim-1).mean() return loss class KLDivLoss(nn.Module): KL 散度损失 def __init__(self, temperature: float 1.0): super().__init__() self.temperature temperature def forward(self, student_logits, teacher_logits): student_log_probs F.log_softmax(student_logits / self.temperature, dim-1) teacher_probs F.softmax(teacher_logits / self.temperature, dim-1) loss F.kl_div(student_log_probs, teacher_probs, reductionbatchmean) return loss * (self.temperature ** 2) class CosineEmbeddingLoss(nn.Module): 余弦嵌入损失 def __init__(self, margin: float 0.0): super().__init__() self.margin margin def forward(self, x1, x2, target): # target1: 相似, target-1: 不相似 cos_sim F.cosine_similarity(x1, x2) loss torch.where( target 1, 1 - cos_sim, torch.clamp(cos_sim - self.margin, min0.0) ) return loss.mean() class HuberLoss(nn.Module): Huber Loss - 对异常值鲁棒 def __init__(self, delta: float 1.0): super().__init__() self.delta delta def forward(self, predictions, targets): diff predictions - targets abs_diff torch.abs(diff) quadratic torch.min(abs_diff, torch.tensor(self.delta)) linear abs_diff - quadratic loss 0.5 * quadratic ** 2 self.delta * linear return loss.mean() # # 组合损失函数 # class CombinedLoss(nn.Module): 组合多种损失函数 def __init__(self, num_classes10, loss_weightsNone): super().__init__() self.ce_loss nn.CrossEntropyLoss() self.focal_loss FocalLoss(alpha0.25, gamma2.0) self.label_smoothing LabelSmoothingLoss(num_classes, smoothing0.1) if loss_weights is None: self.weights {ce: 1.0, focal: 0.5, smoothing: 0.3} else: self.weights loss_weights def forward(self, logits, targets): ce self.ce_loss(logits, targets) focal self.focal_loss(logits, targets) smoothing self.label_smoothing(logits, targets) total (self.weights[ce] * ce self.weights[focal] * focal self.weights[smoothing] * smoothing) components { ce: ce.item(), focal: focal.item(), smoothing: smoothing.item(), total: total.item(), } return total, components # # 梯度检查工具 # class GradientChecker: 检查损失函数的梯度是否正确传播 staticmethod def check_gradient(loss_fn, input_shape, target_shapeNone, devicecpu): 检查梯度是否正确传播 if target_shape is None: target_shape input_shape # 创建随机输入 predictions torch.randn(*input_shape, requires_gradTrue, devicedevice) if isinstance(loss_fn, (TripletLoss, ContrastiveLoss)): # 特殊处理 return GradientChecker._check_triplet(loss_fn, input_shape, device) targets torch.randint(0, input_shape[-1] if len(input_shape) 1 else 10, target_shape[:-1] if len(target_shape) 1 else (target_shape[0],), devicedevice) # 前向传播 loss loss_fn(predictions, targets) # 反向传播 loss.backward() # 检查梯度 has_grad predictions.grad is not None has_nan torch.isnan(predictions.grad).any().item() if has_grad else True has_inf torch.isinf(predictions.grad).any().item() if has_grad else True all_zero (predictions.grad 0).all().item() if has_grad else True print(f损失值: {loss.item():.6f}) print(f梯度存在: {has_grad}) print(f梯度有 NaN: {has_nan}) print(f梯度有 Inf: {has_inf}) print(f梯度全为 0: {all_zero}) if has_grad and not has_nan and not has_inf: print(f梯度范围: [{predictions.grad.min().item():.6f}, f{predictions.grad.max().item():.6f}]) return has_grad and not has_nan and not has_inf and not all_zero staticmethod def _check_triplet(loss_fn, input_shape, device): anchor torch.randn(*input_shape, requires_gradTrue, devicedevice) positive torch.randn(*input_shape, requires_gradTrue, devicedevice) negative torch.randn(*input_shape, requires_gradTrue, devicedevice) loss loss_fn(anchor, positive, negative) loss.backward() has_grad anchor.grad is not None has_nan torch.isnan(anchor.grad).any().item() if has_grad else True all_zero (anchor.grad 0).all().item() if has_grad else True print(f损失值: {loss.item():.6f}) print(f梯度存在: {has_grad}) print(f梯度有 NaN: {has_nan}) print(f梯度全为 0: {all_zero}) return has_grad and not has_nan and not all_zero # # 使用示例 # def demo_focal_loss(): print( * 60) print(示例 1: Focal Loss) print( * 60) loss_fn FocalLoss(alpha0.25, gamma2.0) # 模拟类别不平衡数据 logits torch.randn(100, 10, requires_gradTrue) # 90% 是类别 0 targets torch.cat([torch.zeros(90, dtypetorch.long), torch.randint(1, 10, (10,))]) loss loss_fn(logits, targets) print(fFocal Loss: {loss.item():.4f}) # 对比标准 CE ce nn.CrossEntropyLoss()(logits, targets) print(fStandard CE: {ce.item():.4f}) # 梯度检查 print(\n梯度检查:) GradientChecker.check_gradient(loss_fn, (32, 10)) print() def demo_dice_loss(): print( * 60) print(示例 2: Dice Loss) print( * 60) loss_fn DiceLoss() # 模拟分割预测 predictions torch.sigmoid(torch.randn(4, 1, 32, 32, requires_gradTrue)) targets (torch.rand(4, 1, 32, 32) 0.5).float() loss loss_fn(predictions, targets) print(fDice Loss: {loss.item():.4f}) print(fDice Score: {1 - loss.item():.4f}) print() def demo_triplet_loss(): print( * 60) print(示例 3: Triplet Loss) print( * 60) loss_fn TripletLoss(margin1.0) anchor torch.randn(32, 128, requires_gradTrue) positive torch.randn(32, 128, requires_gradTrue) negative torch.randn(32, 128, requires_gradTrue) loss loss_fn(anchor, positive, negative) print(fTriplet Loss: {loss.item():.4f}) # 梯度检查 print(\n梯度检查:) GradientChecker.check_gradient(loss_fn, (32, 128)) print() def demo_label_smoothing(): print( * 60) print(示例 4: Label Smoothing) print( * 60) loss_fn LabelSmoothingLoss(num_classes5, smoothing0.1) logits torch.randn(4, 5, requires_gradTrue) targets torch.tensor([0, 1, 2, 3]) loss loss_fn(logits, targets) print(fLabel Smoothing Loss: {loss.item():.4f}) # 对比标准 CE ce nn.CrossEntropyLoss()(logits, targets) print(fStandard CE: {ce.item():.4f}) print() def demo_combined_loss(): print( * 60) print(示例 5: 组合损失函数) print( * 60) loss_fn CombinedLoss(num_classes10) logits torch.randn(32, 10, requires_gradTrue) targets torch.randint(0, 10, (32,)) loss, components loss_fn(logits, targets) print(f组合损失: {loss.item():.4f}) print(f各组件: {components}) print() def demo_training_with_custom_loss(): print( * 60) print(示例 6: 使用自定义损失训练) print( * 60) from torch.utils.data import DataLoader, TensorDataset import torch.optim as optim # 创建不平衡数据 torch.manual_seed(42) X torch.randn(500, 10) y torch.cat([torch.zeros(450, dtypetorch.long), torch.ones(50, dtypetorch.long)]) dataset TensorDataset(X, y) dataloader DataLoader(dataset, batch_size32, shuffleTrue) model nn.Sequential( nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 2), ) # 使用 Focal Loss criterion FocalLoss(alpha0.25, gamma2.0) optimizer optim.Adam(model.parameters(), lr0.001) for epoch in range(10): model.train() total_loss 0 correct 0 total 0 for batch_x, batch_y in dataloader: optimizer.zero_grad() output model(batch_x) loss criterion(output, batch_y) loss.backward() optimizer.step() total_loss loss.item() pred output.argmax(dim1) correct pred.eq(batch_y).sum().item() total batch_y.size(0) if (epoch 1) % 5 0: print(fEpoch {epoch1}: Loss{total_loss/len(dataloader):.4f}, fAcc{100.*correct/total:.2f}%) print() def demo_gradient_check(): print( * 60) print(示例 7: 梯度检查) print( * 60) loss_functions [ (FocalLoss, FocalLoss()), (LabelSmoothingLoss, LabelSmoothingLoss(10, 0.1)), (HuberLoss, HuberLoss()), ] for name, loss_fn in loss_functions: print(f\n{name}:) GradientChecker.check_gradient(loss_fn, (32, 10)) print() if __name__ __main__: demo_focal_loss() demo_dice_loss() demo_triplet_loss() demo_label_smoothing() demo_combined_loss() demo_training_with_custom_loss() demo_gradient_check() print( * 60) print(所有示例执行完毕) print( * 60)常见陷阱与注意事项1. 不要在损失函数中使用 numpy# 错误numpy 操作断裂梯度 def wrong_loss(pred, target): pred_np pred.detach().numpy() loss np.mean((pred_np - target.numpy()) ** 2) return torch.tensor(loss) # 正确使用 tensor 操作 def correct_loss(pred, target): return ((pred - target) ** 2).mean()2. 数值稳定性# 错误log(0) -inf loss -torch.log(probabilities) # 正确使用 clamp 或 log_softmax loss -torch.log(probabilities.clamp(min1e-8)) # 或 loss F.nll_loss(F.log_softmax(logits, dim-1), targets)3. 避免不可导操作# 错误argmax 不可导 pred_class predictions.argmax(dim1) loss (pred_class targets).float().mean() # 正确使用可导的替代 loss F.cross_entropy(logits, targets)4. 使用with torch.no_grad()处理不需要梯度的部分class LabelSmoothingLoss(nn.Module): def forward(self, logits, targets): # 创建平滑标签不需要梯度 with torch.no_grad(): true_dist torch.zeros_like(logits) true_dist.fill_(self.smoothing / (self.num_classes - 1)) true_dist.scatter_(1, targets.unsqueeze(1), self.confidence) # 计算损失需要梯度 log_probs F.log_softmax(logits, dim-1) return -(true_dist * log_probs).sum(dim-1).mean()5. 损失函数必须输出标量# 错误返回矩阵 loss (predictions - targets) ** 2 # [batch, classes] # 正确reduction 为标量 loss ((predictions - targets) ** 2).mean() # 标量6. GPU 兼容性class MyLoss(nn.Module): def __init__(self): super().__init__() # 不要在 __init__ 中创建 tensor # self.margin torch.tensor(1.0) # 可能在 CPU 上 def forward(self, x): # 在 forward 中创建 tensor自动在正确设备上 margin torch.tensor(1.0, devicex.device) return F.relu(x - margin).mean()7. 避免在损失函数中使用 Python 循环# 错误极慢 def slow_loss(pred, target): total 0 for i in range(len(pred)): total (pred[i] - target[i]) ** 2 return total / len(pred) # 正确向量化 def fast_loss(pred, target): return ((pred - target) ** 2).mean()总结在 PyTorch 中实现自定义损失函数关键要点如下使用 tensor 操作所有计算必须使用 PyTorch tensor 操作避免 numpy 或 Python 标量确保梯度链不断裂。数值稳定性使用clamp、log_softmax、logsigmoid等函数避免log(0)、exp(overflow)等数值问题。输出标量损失函数必须输出标量值使用mean()或sum()进行 reduction。避免不可导操作不要使用argmax、比较操作等不可导操作使用可导的替代方案。GPU 兼容在forward中创建 tensor 时使用devicex.device确保设备一致。向量化操作避免 Python 循环使用 tensor 的向量化操作提高性能。使用torch.no_grad()对于不需要梯度的部分如创建标签使用no_grad()上下文管理器。梯度检查实现后使用GradientChecker验证梯度是否正确传播确保没有 NaN 或全零梯度。通过遵循这些原则可以实现正确、高效、数值稳定的自定义损失函数满足各种特殊任务的需求。
