以下是基于pytorch的yolov5和arcface人脸识别系统中包含ArcFace函数的models.py文件代码：
```python
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
```
3772
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
