pytorch-loss

我实现了一系列损失函数和操作符，包括标签平滑、AMSoftmax、部分FC、焦点损失、双焦点损失、三元组损失、GIoU/DIoU/CIoU损失及其相关函数、亲和力损失、PC Softmax交叉熵损失、OHEM损失（基于在线困难样本挖掘的Softmax损失）、大 margin Softmax（BMVC 2019）、LoVASZ Softmax损失以及Dice损失（包括广义软Dice损失和批次软Dice损失）。这些可能对我未来的工作有所帮助。

我还尝试实现了Swish、Hard-Swish（HSwish）和Mish激活函数。

此外，新增了一个基于CUDA的一热编码函数（支持标签平滑）。

新添加了一个“指数移动平均（EMA）”算子。

还增加了卷积操作符，例如坐标卷积2D（CoordConv2d）和动态卷积2D（DY_Conv2d）。

部分算子是通过PyTorch的CUDA扩展实现的，因此需要先进行编译：

    $ python -m pip install .

编译安装完成后，您可以根据需要选择并使用相应的损失函数或操作符，例如：

from pytorch_loss import SwishV1, SwishV2, SwishV3
from pytorch_loss import HSwishV1, HSwishV2, HSwishV3
from pytorch_loss import MishV1, MishV2, MishV3
from pytorch_loss import convert_to_one_hot, convert_to_one_hot_cu, OnehotEncoder
from pytorch_loss import EMA

from pytorch_loss import TripletLoss
from pytorch_loss import SoftDiceLossV1, SoftDiceLossV2, SoftDiceLossV3
from pytorch_loss import PCSoftmaxCrossEntropyV1, PCSoftmaxCrossEntropyV2
from pytorch_loss import LargeMarginSoftmaxV1, LargeMarginSoftmaxV2, LargeMarginSoftmaxV3
from pytorch_loss import LabelSmoothSoftmaxCEV1, LabelSmoothSoftmaxCEV2, LabelSmoothSoftmaxCEV3
from pytorch_loss import GIOULoss, DIOULoss, CIOULoss
from pytorch_loss import iou_func, giou_func, diou_func, ciou_func
from pytorch_loss import FocalLossV1, FocalLossV2, FocalLossV3
from pytorch_loss import Dual_Focal_loss
from pytorch_loss import GeneralizedSoftDiceLoss, BatchSoftDiceLoss
from pytorch_loss import AMSoftmax
from pytorch_loss import AffinityFieldLoss, AffinityLoss
from pytorch_loss import OhemCELoss, OhemLargeMarginLoss
from pytorch_loss import LovaszSoftmaxV1, LovaszSoftmaxV3
from pytorch_loss import TaylorCrossEntropyLossV1, TaylorCrossEntropyLossV3
from pytorch_loss import InfoNceDist
from pytorch_loss import PartialFCAMSoftmax

from pytorch_loss import TaylorSoftmaxV1, TaylorSoftmaxV3
from pytorch_loss import LogTaylorSoftmaxV1, LogTaylorSoftmaxV3

from pytorch_loss import CoordConv2d, DY_Conv2d

需要注意的是，某些损失函数或操作符有三个版本，例如LabelSmoothSoftmaxCEV1、LabelSmoothSoftmaxCEV2、LabelSmoothSoftmaxCEV3。其中，“V1”表示仅使用纯PyTorch操作实现，并利用torch.autograd进行反向传播；“V2”同样使用纯PyTorch操作，但采用自行推导的公式进行反向传播；而“V3”则是通过CUDA扩展实现。通常来说，“V3”版本的算子速度更快、内存效率更高，因为我尽量将所有计算压缩到一个CUDA内核函数中，在大多数情况下，这样产生的开销比多个PyTorch操作的组合要小。

如果恰好发现本仓库的用户在代码中发现了错误，请随时提交Issue指出，以便我及时修正。

pytorch-loss 快速上手指南

pytorch-loss 是一个集成了多种常用及前沿损失函数、激活函数和卷积算子的 PyTorch 扩展库。它包含了标签平滑（Label Smooth）、AMSoftmax、Focal Loss、Triplet Loss、GIoU/DIoU/CIoU、Dice Loss 等实现，并提供了基于 CUDA 加速的高性能版本。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux (推荐，因为包含 CUDA 扩展编译) 或 macOS/Windows (部分功能可能受限)。
• Python 版本：建议 Python 3.6+。
• 核心依赖：
  • PyTorch (需安装与您的 CUDA 版本匹配的 PyTorch)
  • torchvision
  • CUDA Toolkit (如需使用 V3 版本的 CUDA 加速算子，必须安装对应的 CUDA 开发包)
• 编译工具：由于部分算子通过 pytorch cuda extension 实现，系统需安装 gcc 或 g++ 编译器。

安装步骤

该库包含自定义 CUDA 算子，因此需要通过源码编译安装。

1. 克隆仓库或下载源码（如果您已有源码包可跳过此步）：

   git clone https://github.com/CoinCheung/pytorch-loss.git
   cd pytorch-loss
   

2. 执行安装命令： 使用 pip 进行本地安装，这将自动触发 CUDA 扩展的编译过程。

   python -m pip install .
   

   提示：如果在国内网络环境下遇到下载依赖慢的问题，建议在 pip 命令后添加国内镜像源参数，例如：

   python -m pip install . -i https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 验证安装： 安装完成后，尝试导入模块确认无报错：

   from pytorch_loss import FocalLossV1
   print("Installation successful!")
   

基本使用

安装成功后，您可以直接从 pytorch-loss 中导入所需的损失函数、激活函数或算子。

1. 导入模块

库中的许多组件提供了三个版本（V1, V2, V3），区别如下：

• V1: 纯 PyTorch 实现，使用 torch.autograd 进行反向传播。
• V2: 纯 PyTorch 实现，但使用自定义推导公式进行反向传播。
• V3: 基于 CUDA 扩展实现，通常速度更快且显存效率更高（推荐在支持 CUDA 的环境中使用）。

# 导入激活函数
from pytorch_loss import SwishV1, HSwishV3, MishV2

# 导入损失函数
from pytorch_loss import FocalLossV3, LabelSmoothSoftmaxCEV3, TripletLoss
from pytorch_loss import GIOULoss, SoftDiceLossV1, AMSoftmax

# 导入工具算子
from pytorch_loss import convert_to_one_hot_cu, EMA
from pytorch_loss import CoordConv2d, DY_Conv2d

2. 使用示例

以下是一个使用 FocalLossV3 和 LabelSmoothSoftmaxCEV3 的简单训练循环片段示例：

import torch
import torch.nn as nn
from pytorch_loss import FocalLossV3, LabelSmoothSoftmaxCEV3

# 初始化模型和优化器 (此处仅为示意)
model = YourModel()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 实例化损失函数
# gamma=2.0, alpha=None 为 Focal Loss 常见默认参数，具体请参考源码或根据任务调整
criterion_focal = FocalLossV3(gamma=2.0) 
criterion_label_smooth = LabelSmoothSoftmaxCEV3(lb_pos=0.9) # lb_pos 控制平滑程度

# 模拟前向传播数据
inputs = torch.randn(4, 10, requires_grad=True).cuda() # Batch=4, Classes=10
targets = torch.tensor([0, 1, 2, 3]).cuda()

# 计算损失
loss_focal = criterion_focal(inputs, targets)
loss_ls = criterion_label_smooth(inputs, targets)

# 反向传播
loss_focal.backward()
# optimizer.step() ...

3. 特殊算子使用示例 (EMA)

使用指数移动平均（EMA）更新模型参数：

from pytorch_loss import EMA

# 初始化 EMA，decay 为衰减率
ema = EMA(model, decay=0.999)

# 在训练步骤结束后更新 EMA 影子模型
ema.update()

# 获取影子模型用于验证或测试
shadow_model = ema.shadow