写一段pytorch人脸识别代码

以下是一个简单的PyTorch人脸识别代码，使用了一个预训练的ResNet50作为特征提取器和分类器：

import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable

# 定义数据变换
transform = transforms.Compose(
    [transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406],
                          std=[0.229, 0.224, 0.225])])

# 加载数据集
trainset = torchvision.datasets.ImageFolder(root='./data/train', transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True, num_workers=2)

# 定义模型
class FaceNet(nn.Module):
    def __init__(self):
        super(FaceNet, self).__init__()
        self.resnet = torchvision.models.resnet50(pretrained=True)
        self.resnet.fc = nn.Linear(2048, 128)

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

# 定义损失函数和优化器
net = FaceNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 训练模型
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs, labels = Variable(inputs), Variable(labels)
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.data[0]
    print('Epoch %d, loss: %.3f' % (epoch+1, running_loss/len(trainloader)))

print('Finished Training')

# 保存模型
torch.save(net.state_dict(), 'facenet.pth')

在训练完成后，可以使用保存的模型进行人脸识别：

import cv2
import numpy as np

# 加载模型
net = FaceNet()
net.load_state_dict(torch.load('facenet.pth'))

# 定义人脸检测器
face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')

# 加载测试图片
img = cv2.imread('test.jpg')

# 转换为灰度图像
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# 检测人脸
faces = face_cascade.detectMultiScale(gray, 1.3, 5)

# 对每个人脸进行识别
for (x,y,w,h) in faces:
    # 提取人脸图像
    face_img = img[y:y+h, x:x+w]
    # 缩放为网络输入大小
    face_img = cv2.resize(face_img, (224, 224))
    # 转换为PyTorch Tensor
    face_tensor = transform(face_img)
    # 添加一个维度以匹配网络输入
    face_tensor = face_tensor.unsqueeze(0)
    # 使用网络进行特征提取
    feature = net(face_tensor)
    # 计算欧几里得距离，找到最相似的人脸
    distances = []
    for i in range(len(features)):
        distance = np.sqrt(np.sum(np.square(features[i].detach().numpy() - feature.detach().numpy())))
        distances.append(distance)
    min_distance = min(distances)
    min_index = distances.index(min_distance)
    # 显示识别结果
    cv2.putText(img, trainset.classes[min_index], (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA)
    cv2.rectangle(img, (x, y), (x+w, y+h), (0, 255, 0), 2)

# 显示结果
cv2.imshow('Image', img)
cv2.waitKey(0)
cv2.destroyAllWindows()

以上代码仅作为示例，实际应用中可能需要更复杂的模型和数据集，并且需要更多的预处理和后处理步骤。