PyTorch 中心损失

PyTorch 中心损失的实现：Wen 等人. 一种用于深度人脸识别的判别特征学习方法. ECCV 2016.

该损失函数也被 deep-person-reid 使用。

快速入门

克隆此仓库并运行代码：

$ git clone https://github.com/KaiyangZhou/pytorch-center-loss
$ cd pytorch-center-loss
$ python main.py --eval-freq 1 --gpu 0 --save-dir log/ --plot

您将在终端中看到以下信息：

当前使用 GPU: 0
创建数据集: mnist
创建模型: cnn
==> 第 1 轮/100 轮
第 50 批/469 批     损失 2.332793 (2.557837) 交叉熵损失 2.332744 (2.388296) 中心损失 0.000048 (0.169540)
第 100 批/469 批    损失 2.354638 (2.463851) 交叉熵损失 2.354637 (2.379078) 中心损失 0.000001 (0.084773)
第 150 批/469 批    损失 2.361732 (2.434477) 交叉熵损失 2.361732 (2.377962) 中心损失 0.000000 (0.056515)
第 200 批/469 批    损失 2.336701 (2.417842) 交叉熵损失 2.336700 (2.375455) 中心损失 0.000001 (0.042386)
第 250 批/469 批    损失 2.404814 (2.407015) 交叉熵损失 2.404813 (2.373106) 中心损失 0.000001 (0.033909)
第 300 批/469 批    损失 2.338753 (2.398546) 交叉熵损失 2.338752 (2.370288) 中心损失 0.000001 (0.028258)
第 350 批/469 批    损失 2.367068 (2.390672) 交叉熵损失 2.367059 (2.366450) 中心损失 0.000009 (0.024221)
第 400 批/469 批    损失 2.344178 (2.384820) 交叉熵损失 2.344142 (2.363620) 中心损失 0.000036 (0.021199)
第 450 批/469 批    损失 2.329708 (2.379460) 交叉熵损失 2.329661 (2.360611) 中心损失 0.000047 (0.018848)
==> 测试
准确率 (%): 10.32  错误率 (%): 89.68
... ...
==> 第 30 轮/100 轮
第 50 批/469 批     损失 0.141117 (0.155986) 交叉熵损失 0.084169 (0.091617) 中心损失 0.056949 (0.064369)
第 100 批/469 批    损失 0.138201 (0.151291) 交叉熵损失 0.089146 (0.092839) 中心损失 0.049055 (0.058452)
第 150 批/469 批    损失 0.151055 (0.151985) 交叉熵损失 0.090816 (0.092405) 中心损失 0.060239 (0.059580)
第 200 批/469 批    损失 0.150803 (0.153333) 交叉熵损失 0.092857 (0.092156) 中心损失 0.057946 (0.061176)
第 250 批/469 批    损失 0.162954 (0.154971) 交叉熵损失 0.094889 (0.092099) 中心损失 0.068065 (0.062872)
第 300 批/469 批    损失 0.162895 (0.156038) 交叉熵损失 0.093100 (0.092034) 中心损失 0.069795 (0.064004)
第 350 批/469 批    损失 0.146187 (0.156491) 交叉熵损失 0.082508 (0.091787) 中心损失 0.063679 (0.064704)
第 400 批/469 批    损失 0.171533 (0.157390) 交叉熵损失 0.092526 (0.091674) 中心损失 0.079007 (0.065716)
第 450 批/469 批    损失 0.209196 (0.158371) 交叉熵损失 0.098388 (0.091560) 中心损失 0.110808 (0.066811)
==> 测试
准确率 (%): 98.51  错误率 (%): 1.49
... ...

有关输入参数的更多详细信息，请运行 python main.py -h。

结果

我们在下方可视化了特征学习过程。

仅使用 Softmax。左：训练集。右：测试集。

Softmax + 中心损失。左：训练集。右：测试集。

如何在您自己的项目中使用中心损失

1. 您只需要 center_loss.py 文件：

from center_loss import CenterLoss

2. 在主函数中初始化中心损失：

center_loss = CenterLoss(num_classes=10, feat_dim=2, use_gpu=True)

3. 为中心损失构建优化器：

optimizer_centloss = torch.optim.SGD(center_loss.parameters(), lr=0.5)

或者，您可以将模型和中心损失的优化器合并，例如：

params = list(model.parameters()) + list(center_loss.parameters())
optimizer = torch.optim.SGD(params, lr=0.1) # 这里 lr 是整体学习率

4. 更新类别中心的方式与更新 PyTorch 模型相同：

# features (torch 张量): 形状为 (batch_size, feat_dim) 的二维浮点张量
# labels (torch 长整型张量): 形状为 (batch_size) 的一维长整型张量
# alpha (float): 中心损失的权重
loss = center_loss(features, labels) * alpha + other_loss
optimizer_centloss.zero_grad()
loss.backward()
# 乘以 (1./alpha)，以消除 alpha 对更新中心的影响
for param in center_loss.parameters():
    param.grad.data *= (1./alpha)
optimizer_centloss.step()

如果您采用第二种方式（即使用一个优化器同时优化模型和中心损失），更新代码将如下所示：

loss = center_loss(features, labels) * alpha + other_loss
optimizer.zero_grad()
loss.backward()
for param in center_loss.parameters():
    # lr_cent 是中心损失的学习率，例如 lr_cent = 0.5
    param.grad.data *= (lr_cent / (alpha * lr))
optimizer.step()

pytorch-center-loss 快速上手指南

环境准备

• 系统要求：Linux / macOS / Windows
• 前置依赖：
  • Python 3.x
  • PyTorch (建议版本 >= 1.0)
  • torchvision
  • matplotlib (用于绘图展示)

国内加速建议：推荐使用清华或阿里镜像源安装依赖，以提升下载速度。

pip install torch torchvision -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

克隆仓库并进入目录：

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

(注：若 GitHub 访问缓慢，可尝试使用国内代码托管平台镜像或配置代理)

基本使用

1. 运行示例代码

直接运行提供的 main.py 脚本，在 MNIST 数据集上演示 Center Loss 的训练过程：

python main.py --eval-freq 1 --gpu 0 --save-dir log/ --plot

运行后终端将输出每个 Epoch 的 Loss 变化（包含 XentLoss 和 CenterLoss）以及测试集准确率。

2. 集成到自己的项目

只需引入 center_loss.py 文件即可使用。

第一步：导入模块

from center_loss import CenterLoss

第二步：初始化损失函数

# num_classes: 类别数量, feat_dim: 特征维度, use_gpu: 是否使用 GPU
center_loss = CenterLoss(num_classes=10, feat_dim=2, use_gpu=True)

第三步：构建优化器 你可以为 Center Loss 单独创建优化器，也可以与模型参数合并：

# 方式 A：单独优化器
optimizer_centloss = torch.optim.SGD(center_loss.parameters(), lr=0.5)

# 方式 B：合并优化器 (推荐)
params = list(model.parameters()) + list(center_loss.parameters())
optimizer = torch.optim.SGD(params, lr=0.1)

第四步：计算损失并更新参数 在训练循环中，注意需要对 Center Loss 的梯度进行缩放处理，以消除权重系数 alpha 对中心点更新的影响。

若使用单独优化器：

alpha = 0.5  # center loss 的权重
features = model(input_data)  # 形状：(batch_size, feat_dim)
labels = target_data          # 形状：(batch_size)

loss = center_loss(features, labels) * alpha + cross_entropy_loss
optimizer_centloss.zero_grad()
loss.backward()

# 关键步骤：抵消 alpha 对中心点梯度的影响
for param in center_loss.parameters():
    param.grad.data *= (1./alpha)

optimizer_centloss.step()
optimizer_model.step() # 更新主模型

若使用合并优化器：

alpha = 0.5
lr_cent = 0.5  # center loss 的学习率
lr = 0.1       # 整体学习率

loss = center_loss(features, labels) * alpha + cross_entropy_loss
optimizer.zero_grad()
loss.backward()

# 关键步骤：根据学习率比例调整梯度
for param in center_loss.parameters():
    param.grad.data *= (lr_cent / (alpha * lr))

optimizer.step()