Centerloss的pytorch实现

以下是使用PyTorch实现Center Loss的代码示例：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

class CenterLoss(nn.Module):
    def __init__(self, num_classes, feat_dim, loss_weight=0.5):
        super(CenterLoss, self).__init__()
        self.num_classes = num_classes
        self.feat_dim = feat_dim
        self.loss_weight = loss_weight
        self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim))

    def forward(self, x, labels):
        batch_size = x.size(0)
        # 计算当前batch中每个样本对应的中心
        centers_batch = self.centers[labels]
        # 计算当前batch中每个样本与其对应中心之间的距离
        dist = torch.sum((x - centers_batch) ** 2, dim=1)
        # 计算center loss
        center_loss = torch.mean(dist)

        # 更新中心
        diff = centers_batch - x
        unique_label, unique_idx = torch.unique(labels, return_inverse=True)
        appear_times = torch.bincount(unique_idx, minlength=self.num_classes).float()
        appear_times = appear_times.unsqueeze(1)
        centers_update = torch.zeros_like(self.centers)
        centers_update.scatter_add_(0, labels.view(-1, 1).expand(batch_size, self.feat_dim), diff)
        centers_update = centers_update / (appear_times + 1e-8)
        self.centers.data = self.centers.data - self.loss_weight * centers_update.data

        return center_loss

class Net(nn.Module):
    def __init__(self, num_classes, feat_dim):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(feat_dim, 512)
        self.fc2 = nn.Linear(512, num_classes)
        self.center_loss = CenterLoss(num_classes, feat_dim)

    def forward(self, x, labels):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        center_loss = self.center_loss(x, labels)
        return x, center_loss

在这里，我们首先定义了一个`CenterLoss`类来计算中心损失。`CenterLoss`的`__init__`函数中包含中心矩阵`centers`，其大小为`(num_classes, feat_dim)`，其中`num_classes`为类别数，`feat_dim`为特征维度。`forward`函数接受输入的特征张量`x`和对应的标签`labels`，计算`x`和每个样本对应的中心之间的距离，然后计算中心损失并更新中心矩阵。在`Net`类中，我们将`CenterLoss`作为一个模块集成到模型中，同时在模型的前向传播中计算中心损失。

接下来，我们可以使用以下代码来训练模型：

net = Net(num_classes, feat_dim)
optimizer = optim.SGD(net.parameters(), lr=lr, momentum=momentum)

for epoch in range(num_epochs):
    for batch_idx, (data, labels) in enumerate(train_loader):
        data, labels = data.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs, center_loss = net(data, labels)
        softmax_loss = F.cross_entropy(outputs, labels)
        loss = softmax_loss + center_loss
        loss.backward()
        optimizer.step()

在每个batch的训练中，我们首先将输入数据和标签送入设备中，然后将模型参数的梯度清零。接着，我们计算前向传播的结果和中心损失，然后使用交叉熵损失计算总损失，并进行反向传播和参数更新。