Pytorch的中心损失怎么使用，给一个具体的案例

PyTorch的中心损失可以用于增强特征表示学习。 以下是使用中心损失的一个例子：

import torch
import torch.nn as nn
import torch.optim as optim

class CustomLoss(nn.Module):
    def __init__(self, num_classes, feat_dim, size_average=True):
        super(CustomLoss, self).__init__()
        self.num_classes = num_classes
        self.feat_dim = feat_dim
        self.size_average = size_average
        self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim).cuda()) # 初始化中心

    def forward(self, inputs, targets):
        """
        Args:
            inputs: feature matrix with shape (batch_size, feat_dim).
            targets: ground truth labels with shape (num_classes).
        """
        batch_size = inputs.size(0)
        dif = inputs.unsqueeze(1) - self.centers.unsqueeze(0)  # 计算 features 和 centers 之间的差异
        l2_dists = torch.sqrt(torch.sum(torch.pow(dif, 2), 2))  # 计算每个 feature 与所有 center 之间的距离
        classes = torch.arange(self.num_classes).long().cuda()
        targets = targets.unsqueeze(1).expand(batch_size, self.num_classes)
        masks = torch.eq(targets, classes.expand(batch_size, self.num_classes))  # 生成掩码
        f_m = l2_dists * masks.float() # 过滤掩码内的距离
        f_m=f_m.cuda()
        loss = f_m.sum() / (self.size_average * batch_size)  # 计算平均损失
        return loss

# 定义模型
class Model(nn.Module):
    def __init__(self, num_classes):
        super(Model, self).__init__()
        self.fc1 = nn.Linear(10, 100)
        self.relu = nn.ReLU(inplace=True)
        self.fc2 = nn.Linear(100, num_classes)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# 定义优化器
model = Model(num_classes=10).cuda()
criterion_cls = nn.CrossEntropyLoss().cuda()
criterion_center = CustomLoss(num_classes=10, feat_dim=100).cuda()
optimizer = optim.SGD([{'params': model.parameters()}, {'params': criterion_center.parameters()}], lr=0.1, momentum=0.5)

# 训练模型
for i in range(100):
    inputs = torch.randn(64, 10).cuda()
    targets = torch.randint(0, 10, (64,)).cuda()

    optimizer.zero_grad()  # 梯度清零

    outputs = model(inputs)  # 前向传播
    loss_cls = criterion_cls(outputs, targets)  # 分类损失
    loss_center = criterion_center(model.fc1.weight, targets)  # 中心损失

    loss = loss_cls + loss_center  # 总损失
    loss.backward()  # 反向传播

    optimizer.step()  # 更新参数

在此示例中，我们定义了一个自定义损失函数 `CustomLoss`，其中包含一个参数 `centres`，这些参数是特征中心，它们在每次反向传播时进行更新。我们还定义了一个模型 `Model`，它包含一个线性层和一个ReLU激活函数。在训练循环中，我们计算了分类损失和中心损失，并将它们相加作为总损失。我们使用 `torch.optim.SGD` 优化器进行参数更新。