基于pytorch的yolov5和arcface人脸识别系统中中包含ArcFace函数的models.py文件代码

以下是基于pytorch的yolov5和arcface人脸识别系统中包含ArcFace函数的models.py文件代码：

import torch
import torch.nn as nn
import math
import torch.nn.functional as F

from collections import OrderedDict

# Convolutional block
class Conv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, groups=1, bias=False, dilation=1):
        super(Conv, self).__init__()
        padding = (kernel_size - 1) // 2 * dilation
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, groups=groups, bias=bias, dilation=dilation)
        self.bn = nn.BatchNorm2d(out_channels)
        self.act = nn.LeakyReLU(0.1)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.act(x)
        return x

# Focus wh information into c space
class Focus(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, groups=1, bias=False, dilation=1):
        super(Focus, self).__init__()
        self.conv = Conv(in_channels * 4, out_channels, kernel_size, stride, groups, bias, dilation)

    def forward(self, x):
        # Concatenate the patches by channel and spatially
        x = torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], dim=1)
        x = self.conv(x)
        return x

# CSP Bottleneck with 3 convolutions
class BottleneckCSP(nn.Module):
    def __init__(self, in_channels, out_channels, shortcut=True, groups=1, expansion=0.5):
        super(BottleneckCSP, self).__init__()
        hidden_channels = int(out_channels * expansion)
        self.conv1 = Conv(in_channels, hidden_channels, 1)
        self.conv2 = Conv(in_channels, hidden_channels, 1)
        self.conv3 = Conv(hidden_channels, out_channels, 1)
        self.conv4 = Conv(hidden_channels, out_channels, 3, groups=groups)
        self.bn = nn.BatchNorm2d(2 * out_channels)
        self.act = nn.LeakyReLU(0.1)
        self.shortcut = shortcut

    def forward(self, x):
        shortcut = x
        x1 = self.conv1(x)
        x2 = self.conv2(x)
        x2 = self.conv4(x2)
        x = torch.cat([x1, x2], dim=1)
        x = self.conv3(x)
        x = self.bn(x)
        if self.shortcut:
            x = shortcut + x
        x = self.act(x)
        return x

# SPP block
class SPP(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_sizes=[5, 9, 13]):
        super(SPP, self).__init__()
        hidden_channels = in_channels // 2
        self.conv1 = Conv(in_channels, hidden_channels, 1)
        self.conv2 = Conv(hidden_channels * (len(kernel_sizes) + 1), out_channels, 1)
        self.maxpools = nn.ModuleList()
        for kernel_size in kernel_sizes:
            self.maxpools.append(nn.MaxPool2d(kernel_size, stride=1, padding=kernel_size // 2))

    def forward(self, x):
        x1 = self.conv1(x)
        x2s = [x1]
        for maxpool in self.maxpools:
            x2 = maxpool(x)
            x2s.append(x2)
        x = torch.cat(x2s, dim=1)
        x = self.conv2(x)
        return x

# CSP Bottleneck with SPP
class BottleneckCSPSPP(nn.Module):
    def __init__(self, in_channels, out_channels, shortcut=True, groups=1, expansion=0.5):
        super(BottleneckCSPSPP, self).__init__()
        hidden_channels = int(out_channels * expansion)
        self.conv1 = Conv(in_channels, hidden_channels, 1)
        self.conv2 = Conv(in_channels, hidden_channels, 1)
        self.conv3 = Conv(hidden_channels, out_channels, 1)
        self.conv4 = SPP(hidden_channels, hidden_channels)
        self.conv5 = Conv(hidden_channels, out_channels, 1)
        self.bn = nn.BatchNorm2d(2 * out_channels)
        self.act = nn.LeakyReLU(0.1)
        self.shortcut = shortcut

    def forward(self, x):
        shortcut = x
        x1 = self.conv1(x)
        x2 = self.conv2(x)
        x2 = self.conv4(x2)
        x3 = self.conv5(x2)
        x = torch.cat([x1, x3], dim=1)
        x = self.conv3(x)
        x = self.bn(x)
        if self.shortcut:
            x = shortcut + x
        x = self.act(x)
        return x

# CSP Darknet with SPP and PAN
class CSPDarknetSPP(nn.Module):
    def __init__(self, num_classes=1000, width=1.0):
        super(CSPDarknetSPP, self).__init__()
        self.stem = Focus(3, int(64 * width))
        self.layer1 = nn.Sequential(OrderedDict([
            ('bottleneck0', BottleneckCSP(int(64 * width), int(128 * width))),
            ('bottleneck1', BottleneckCSP(int(128 * width), int(128 * width), shortcut=False)),
            ('bottleneck2', BottleneckCSP(int(128 * width), int(128 * width), shortcut=False)),
            ('bottleneck3', BottleneckCSP(int(128 * width), int(128 * width), shortcut=False)),
        ]))
        self.layer2 = nn.Sequential(OrderedDict([
            ('bottleneck4', BottleneckCSP(int(128 * width), int(256 * width))),
            ('bottleneck5', BottleneckCSP(int(256 * width), int(256 * width), shortcut=False)),
            ('bottleneck6', BottleneckCSP(int(256 * width), int(256 * width), shortcut=False)),
            ('bottleneck7', BottleneckCSP(int(256 * width), int(256 * width), shortcut=False)),
            ('bottleneck8', BottleneckCSP(int(256 * width), int(256 * width), shortcut=False)),
            ('bottleneck9', BottleneckCSP(int(256 * width), int(256 * width), shortcut=False)),
            ('bottleneck10', BottleneckCSP(int(256 * width), int(256 * width), shortcut=False)),
            ('bottleneck11', BottleneckCSP(int(256 * width), int(256 * width), shortcut=False)),
        ]))
        self.layer3 = nn.Sequential(OrderedDict([
            ('bottleneck12', BottleneckCSP(int(256 * width), int(512 * width))),
            ('bottleneck13', BottleneckCSP(int(512 * width), int(512 * width), shortcut=False)),
            ('bottleneck14', BottleneckCSP(int(512 * width), int(512 * width), shortcut=False)),
            ('bottleneck15', BottleneckCSP(int(512 * width), int(512 * width), shortcut=False)),
            ('bottleneck16', BottleneckCSP(int(512 * width), int(512 * width), shortcut=False)),
            ('bottleneck17', BottleneckCSP(int(512 * width), int(512 * width), shortcut=False)),
            ('bottleneck18', BottleneckCSP(int(512 * width), int(512 * width), shortcut=False)),
            ('bottleneck19', BottleneckCSP(int(512 * width), int(512 * width), shortcut=False)),
        ]))
        self.layer4 = nn.Sequential(OrderedDict([
            ('bottleneck20', BottleneckCSP(int(512 * width), int(1024 * width))),
            ('bottleneck21', BottleneckCSP(int(1024 * width), int(1024 * width), shortcut=False)),
        ]))
        self.conv = Conv(int(1024 * width), int(512 * width), 1)
        self.bn = nn.BatchNorm2d(int(512 * width))
        self.act = nn.LeakyReLU(0.1)
        self.arcface = ArcFace(int(512 * width), num_classes)

    def forward(self, x):
        x = self.stem(x)
        x = F.max_pool2d(x, kernel_size=2, stride=2)
        x = self.layer1(x)
        x = F.max_pool2d(x, kernel_size=2, stride=2)
        x = self.layer2(x)
        x3 = F.max_pool2d(x, kernel_size=2, stride=2)
        x = self.layer3(x3)
        x4 = F.max_pool2d(x, kernel_size=2, stride=2)
        x = self.layer4(x4)
        x = self.conv(x)
        x = self.bn(x)
        x = self.act(x)
        x = F.adaptive_avg_pool2d(x, (1, 1)).squeeze(-1).squeeze(-1)
        x = self.arcface(x)
        return x

# ArcFace head
class ArcFace(nn.Module):
    def __init__(self, in_features, out_features):
        super(ArcFace, self).__init__()
        self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features))
        self.weight.data.normal_(0, 0.01)
        self.margin = nn.Parameter(torch.FloatTensor([0.5]))
        self.margin.requiresGrad = False
        self.cos_m = math.cos(self.margin)
        self.sin_m = math.sin(self.margin)
        self.mm = self.sin_m * self.margin
        self.threshold = math.cos(math.pi - self.margin)

    def forward(self, x):
        cosine = F.linear(F.normalize(x), F.normalize(self.weight))
        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
        phi = cosine * self.cos_m - sine * self.sin_m
        phi = torch.where(cosine > self.threshold, phi, cosine - self.mm)
        return phi