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)