DeepLearing—CV系列（二十一）——基于arcloss实现的人脸识别项目的Pytorch实现

代码目录：

关于项目流程，可以跳转我的这篇博文查看DeepLearing—CV系列（二十）——基于MTCNN与centerloss/arcloss实现的人脸识别项目流程，本篇博文在其基础上用Pytorch实现了人脸识别项目。

一、face.py

创建arcsoftmax分类器，创建人脸特征提取器（使用预训练模型）

import torchvision.models as models
import torch.nn as nn
import torch
import torch.nn.functional as F

class Arcsoftmax(nn.Module):
    def __init__(self, feature_num, cls_num):
        super().__init__()
        self.W = nn.Parameter(torch.randn((feature_num, cls_num)), requires_grad=True)
        self.func = nn.Softmax()

    def forward(self, x, s=1, m=0.2):
        x_norm = F.normalize(x, dim=1)
        w_norm = F.normalize(self.W, dim=0)

        cosa = torch.matmul(x_norm, w_norm) / 10
        a = torch.acos(cosa)

        arcsoftmax = torch.exp(
            s * torch.cos(a + m) * 10) / (torch.sum(torch.exp(s * cosa * 10), dim=1, keepdim=True) - torch.exp(
            s * cosa * 10) + torch.exp(s * torch.cos(a + m) * 10))

        return arcsoftmax

class FaceNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.sub_net = nn.Sequential(
            models.mobilenet_v2(pretrained=True),
        )
        self.feature_net = nn.Sequential(
            nn.BatchNorm1d(1000),
            nn.LeakyReLU(0.1),
            nn.Linear(1000, 512, bias=False),
        )
        self.arc_softmax = Arcsoftmax(512, 26)

    def forward(self,x):
        y = self.sub_net(x)
        feature = self.feature_net(y)
        return feature,self.arc_softmax(feature,1,1)

    def encode(self,x):
        return self.feature_net(self.sub_net(x))

def compare(face1, face2):
    face1_norm = F.normalize(face1)
    face2_norm = F.normalize(face2)
    print(face1_norm.shape)
    print(face2_norm.shape)
    cosa = torch.matmul(face1_norm, face2_norm.T)
    return cosa

二、Mydataset.py

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os

tf = transforms.Compose([
    transforms.Resize([112, 112]),
    transforms.ToTensor()
])


class MyDataset(Dataset):
    def __init__(self, main_dir):
        self.dataset = []
        for face_dir in os.listdir(main_dir):
            for face_filename in os.listdir(os.path.join(main_dir, face_dir)):
                self.dataset.append([os.path.join(main_dir, face_dir, face_filename), int(face_dir)])

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, item):
        data = self.dataset[item]
        image_data = tf(Image.open(data[0]))
        image_label = data[1]
        return image_data, image_label

if __name__ == '__main__':
    import torch
    mydataset = MyDataset("face_data")
    dataset = DataLoader(mydataset,100,shuffle=True)
    for data in dataset:

        print(data[0].shape)
        print(data[1].shape)
        print(len(data[1]))

三、Train.py

from face import *
import os
from Mydataset import MyDataset
from torch.utils import data
class Trainer:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.net = FaceNet().to(self.device)

    def train(self):
        save_path = "params/1.pt"
        if not os.path.exists("params"):
            os.mkdir("params")
        if os.path.exists(save_path):
            self.net.load_state_dict(torch.load(save_path))
        loss_fn = nn.NLLLoss()
        opt = torch.optim.Adam(self.net.parameters())
        mydataset = MyDataset("face_data")
        dataloader = data.DataLoader(dataset=mydataset,shuffle=True,batch_size=100)

        loss = 0
        for epochs in range(10000):
            for xs,ys in dataloader:
                xs = xs.to(self.device)
                ys = ys.to(self.device)
                feature,cls = self.net(xs)
                loss = loss_fn(torch.log(cls), ys)

                opt.zero_grad()
                loss.backward()
                opt.step()

                predict = torch.argmax(cls,dim=1)
                label = ys
            print("predict:{}\nlabel:{}".format(predict, label))
            accuracy = torch.mean((predict == label), dtype=torch.float32)
            print("accuracy:{}".format(str(round(accuracy.item() * 100))) + "%")
            print(str([epochs]) + "Loss:" + str(loss.item()))

            if epochs % 100 == 0:
                torch.save(self.net.state_dict(), save_path)
                print(str(epochs) + "参数保存成功")

if __name__ == '__main__':

    t = Trainer()
    t.train()

四、use.py

from PIL import ImageDraw,ImageFont,Image
from face import *
from Mydataset import tf

class using:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.save_path = "params/1.pt"
        self.net = FaceNet().to(self.device)
        self.net.load_state_dict(torch.load(self.save_path))
        self.net.eval()

    def us(self):
        person1 = tf(Image.open("test_img/pic6_0.jpg")).to(self.device)
        person1_feature = self.net.encode(torch.unsqueeze(person1,0))
        # person1_feature = net.encode(person1[None, ...])
        # print(person1.shape)
        # print(torch.unsqueeze(person1, 0).shape)
        # print(person1[None, ...].shape)
        person2 = tf(Image.open("test_img/pic5_1.jpg")).to(self.device)
        person2_feature = self.net.encode(person2[None, ...])
        siam = compare(person1_feature, person2_feature)
        print(siam)
        x = "周杰伦" if round(siam.item()) == 1 else "其他人"
        font = ImageFont.truetype("simhei.ttf", 20)
        with Image.open("test_img/pic5_1.jpg") as img:
            imgdraw = ImageDraw.Draw(img)
            imgdrawa = imgdraw.text((0, 0), x, font=font)
            img.show(imgdrawa)
        print()

if __name__ == '__main__':
    u = using()
    u.us()
    # 把模型和参数进行打包，以便C++或PYTHON调用
    # import torch.jit as jit
    # x = torch.Tensor(1, 3, 112, 112)
    # net = FaceNet()
    # net.load_state_dict(torch.load("params/1.pt"))
    # net.eval()
    # traced_script_module = jit.trace(net, x)
    # traced_script_module.save("model.cpt")

五、效果

这里传一下数据集大家可以简单测试一下：
链接：https://pan.baidu.com/s/1q9f-UK6t2fvRmzJHcZZPCA
提取码：kdd4