看到 center loss 可以像聚类一样,使用不同的核函数(可以看成是计算距离的函数)计算loss,记录一下这个demo。
import os
from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
import torch
from torch import nn
class CenterLoss(nn.Module):
"""Center loss.
Reference:
Wen et al. A Discriminative Feature Learning Approach for Deep Face Recognition. ECCV 2016.
Args:
num_classes (int): number of classes.
feat_dim (int): feature dimension.
"""
def __init__(self, num_classes, feat_dim):
super(CenterLoss, self).__init__()
self.num_classes = num_classes
self.feat_dim = feat_dim
self.centers = nn.Parameter(torch.randn(self.num_classes, self.feat_dim))
def forward(self, x, labels):
"""
Args:
x: feature matrix with shape (batch_size, feat_dim).
labels: ground truth labels with shape (batch_size).
"""
device = x.device
batch_size = x.size(0)
if device != self.centers.device:
self.centers = self.centers.to(device)
distmat = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(batch_size, self.num_classes) + \
torch.pow(self.centers, 2).sum(dim=1, keepdim=True).expand(self.num_classes, batch_size).t()
distmat.addmm_(1, -2, x, self.centers.t())
classes = torch.arange(self.num_classes).long().to(device)
labels = labels.unsqueeze(1).expand(batch_size, self.num_classes)
mask = labels.eq(classes.expand(batch_size, self.num_classes))
dist = distmat * mask.float()
loss = dist.clamp(min=0, max=1e12).sum() / batch_size
return loss
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
os.chdir(os.path.dirname(__file__) + "/../")
cLoss = CenterLoss(10, 128)
mnist_train = datasets.MNIST(root='./data/', train=True, transform=transforms.ToTensor(), download=True)
mnist_test = datasets.MNIST(root='./data/', train=False, transform=transforms.ToTensor(), download=True)
net = nn.Sequential(
nn.Linear(28*28, 256),
nn.ReLU(),
nn.Linear(256, 128),
nn.ReLU(),
nn.Linear(128, 128)
)
net = net.cuda()
cLoss = cLoss.cuda()
optimizer = torch.optim.SGD(net.parameters(), lr=0.01)
optimizer_center = torch.optim.SGD(cLoss.parameters(), lr=0.5)
train_loader = DataLoader(mnist_train, batch_size=64, shuffle=True)
test_loader = DataLoader(mnist_test, batch_size=64, shuffle=False)
for epoch in range(50):
net.train()
cLoss.train()
for batch_idx, (data, target) in enumerate(train_loader):
data = data.view(-1, 28*28).cuda()
target = target.cuda()
feature = net(data)
loss1 = nn.CrossEntropyLoss()(feature, target)
loss2 = cLoss(feature, target)
loss = loss1 + 0.1 * loss2
optimizer.zero_grad()
optimizer_center.zero_grad()
loss.backward()
optimizer.step()
optimizer_center.step()
if batch_idx % 100 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
net.eval()
cLoss.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data = data.view(-1, 28*28).cuda()
target = target.cuda()
feature = net(data)
test_loss += nn.CrossEntropyLoss()(feature, target).item()
pred = feature.max(1, keepdim=True)[1]
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))