125 lines
3.9 KiB
Python
125 lines
3.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
|
|
class Flatten(nn.Module):
|
|
def forward(self, x):
|
|
return x.view(x.shape[0], -1)
|
|
|
|
class ConvBn(nn.Module):
|
|
|
|
def __init__(self, in_c, out_c, kernel=(1, 1), stride=1, padding=0, groups=1):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
nn.Conv2d(in_c, out_c, kernel, stride, padding, groups=groups, bias=False),
|
|
nn.BatchNorm2d(out_c)
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
|
|
class ConvBnPrelu(nn.Module):
|
|
|
|
def __init__(self, in_c, out_c, kernel=(1, 1), stride=1, padding=0, groups=1):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
ConvBn(in_c, out_c, kernel, stride, padding, groups),
|
|
nn.PReLU(out_c)
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
|
|
class DepthWise(nn.Module):
|
|
|
|
def __init__(self, in_c, out_c, kernel=(3, 3), stride=2, padding=1, groups=1):
|
|
super().__init__()
|
|
self.net = nn.Sequential(
|
|
ConvBnPrelu(in_c, groups, kernel=(1, 1), stride=1, padding=0),
|
|
ConvBnPrelu(groups, groups, kernel=kernel, stride=stride, padding=padding, groups=groups),
|
|
ConvBn(groups, out_c, kernel=(1, 1), stride=1, padding=0),
|
|
)
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
|
|
class DepthWiseRes(nn.Module):
|
|
"""DepthWise with Residual"""
|
|
|
|
def __init__(self, in_c, out_c, kernel=(3, 3), stride=2, padding=1, groups=1):
|
|
super().__init__()
|
|
self.net = DepthWise(in_c, out_c, kernel, stride, padding, groups)
|
|
|
|
def forward(self, x):
|
|
return self.net(x) + x
|
|
|
|
|
|
class MultiDepthWiseRes(nn.Module):
|
|
|
|
def __init__(self, num_block, channels, kernel=(3, 3), stride=1, padding=1, groups=1):
|
|
super().__init__()
|
|
|
|
self.net = nn.Sequential(*[
|
|
DepthWiseRes(channels, channels, kernel, stride, padding, groups)
|
|
for _ in range(num_block)
|
|
])
|
|
|
|
def forward(self, x):
|
|
return self.net(x)
|
|
|
|
|
|
class FaceMobileNet(nn.Module):
|
|
|
|
def __init__(self, embedding_size):
|
|
super().__init__()
|
|
self.conv1 = ConvBnPrelu(1, 64, kernel=(3, 3), stride=2, padding=1)
|
|
self.conv2 = ConvBn(64, 64, kernel=(3, 3), stride=1, padding=1, groups=64)
|
|
self.conv3 = DepthWise(64, 64, kernel=(3, 3), stride=2, padding=1, groups=128)
|
|
self.conv4 = MultiDepthWiseRes(num_block=4, channels=64, kernel=3, stride=1, padding=1, groups=128)
|
|
self.conv5 = DepthWise(64, 128, kernel=(3, 3), stride=2, padding=1, groups=256)
|
|
self.conv6 = MultiDepthWiseRes(num_block=6, channels=128, kernel=(3, 3), stride=1, padding=1, groups=256)
|
|
self.conv7 = DepthWise(128, 128, kernel=(3, 3), stride=2, padding=1, groups=512)
|
|
self.conv8 = MultiDepthWiseRes(num_block=2, channels=128, kernel=(3, 3), stride=1, padding=1, groups=256)
|
|
self.conv9 = ConvBnPrelu(128, 512, kernel=(1, 1))
|
|
self.conv10 = ConvBn(512, 512, groups=512, kernel=(7, 7))
|
|
self.flatten = Flatten()
|
|
self.linear = nn.Linear(2048, embedding_size, bias=False)
|
|
self.bn = nn.BatchNorm1d(embedding_size)
|
|
|
|
def forward(self, x):
|
|
#print('x',x.shape)
|
|
out = self.conv1(x)
|
|
out = self.conv2(out)
|
|
out = self.conv3(out)
|
|
out = self.conv4(out)
|
|
out = self.conv5(out)
|
|
out = self.conv6(out)
|
|
out = self.conv7(out)
|
|
out = self.conv8(out)
|
|
out = self.conv9(out)
|
|
out = self.conv10(out)
|
|
out = self.flatten(out)
|
|
out = self.linear(out)
|
|
out = self.bn(out)
|
|
return out
|
|
|
|
if __name__ == "__main__":
|
|
from PIL import Image
|
|
import numpy as np
|
|
|
|
x = Image.open("../samples/009.jpg").convert('L')
|
|
x = x.resize((128, 128))
|
|
x = np.asarray(x, dtype=np.float32)
|
|
x = x[None, None, ...]
|
|
x = torch.from_numpy(x)
|
|
net = FaceMobileNet(512)
|
|
net.eval()
|
|
with torch.no_grad():
|
|
out = net(x)
|
|
print(out.shape)
|