基于pytorch环境的arcface训练模型并测试模型的代码

以下是基于 PyTorch 环境的 ArcFace 训练和测试模型的代码：

import torch
import torch.nn as nn
from torchvision import transforms, datasets
import torch.utils.data as data
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.autograd import Variable
import numpy as np
import time
from sklearn.metrics import accuracy_score

# 定义 ArcFace 损失函数
class ArcfaceLoss(nn.Module):
    def __init__(self, s=64.0, m=0.5):
        super(ArcfaceLoss, self).__init__()
        self.s = s
        self.m = m

    def forward(self, inputs, labels):
        cos_theta = inputs
        cos_theta = torch.clamp(cos_theta, -1, 1)
        phi = cos_theta - self.m
        one_hot = torch.zeros(cos_theta.size()).cuda()
        one_hot.scatter_(1, labels.view(-1, 1).long(), 1)
        output = (one_hot * phi) + ((1.0 - one_hot) * cos_theta)
        output *= self.s
        return output


# 定义 ResNet50 网络
class ResNet50(nn.Module):
    def __init__(self, num_classes=10):
        super(ResNet50, self).__init__()
        self.resnet50 = torch.hub.load('pytorch/vision:v0.6.0', 'resnet50', pretrained=True)
        self.resnet50.fc = nn.Linear(2048, num_classes)

    def forward(self, x):
        x = self.resnet50(x)
        return x


# 定义数据增强方法
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

test_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 加载数据集
train_dataset = datasets.ImageFolder(root='train', transform=train_transform)
test_dataset = datasets.ImageFolder(root='test', transform=test_transform)

# 定义数据加载器
train_loader = data.DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
test_loader = data.DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)

# 定义模型和损失函数
model = ResNet50(num_classes=10).cuda()
criterion = ArcfaceLoss().cuda()

# 定义优化器和学习率调整方法
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
scheduler = lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)

# 训练模型
def train(model, criterion, optimizer, scheduler, num_epochs=10):
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch+1, num_epochs))
        print('-' * 10)
        model.train()
        running_loss = 0.0
        for i, (inputs, labels) in enumerate(train_loader):
            inputs = inputs.cuda()
            labels = labels.cuda()
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item() * inputs.size(0)
            if i % 10 == 9:
                print('Batch {}/{} Loss: {:.4f}'.format(i+1, len(train_loader), loss.item()))
        scheduler.step()
        epoch_loss = running_loss / len(train_dataset)
        print('Train Loss: {:.4f}'.format(epoch_loss))
        test(model)

# 测试模型
def test(model):
    model.eval()
    running_corrects = 0
    for inputs, labels in test_loader:
        inputs = inputs.cuda()
        labels = labels.cuda()
        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        running_corrects += torch.sum(preds == labels.data)
    accuracy = running_corrects.double() / len(test_dataset)
    print('Test Accuracy: {:.4f}'.format(accuracy))

if __name__ == '__main__':
    train(model, criterion, optimizer, scheduler, num_epochs=10)

在上述代码中，我们使用 PyTorch 内置的 ResNet-50 网络，并将其用作 ArcFace 训练模型的基础。我们还定义了 ArcFace 损失函数，并使用 SGD 优化器和学习率调整方法对模型进行训练。最后，我们使用测试数据集对模型进行测试，并计算测试精度。