RabbitRUI/ruispace

1'''2Adapted from https://github.com/cavalleria/cavaface.pytorch/blob/master/backbone/mobilefacenet.py3Original author cavalleria4'''5 6import torch.nn as nn7from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, Sequential, Module8import torch9 10 11class Flatten(Module):12    def forward(self, x):13        return x.view(x.size(0), -1)14 15 16class ConvBlock(Module):17    def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):18        super(ConvBlock, self).__init__()19        self.layers = nn.Sequential(20            Conv2d(in_c, out_c, kernel, groups=groups, stride=stride, padding=padding, bias=False),21            BatchNorm2d(num_features=out_c),22            PReLU(num_parameters=out_c)23        )24 25    def forward(self, x):26        return self.layers(x)27 28 29class LinearBlock(Module):30    def __init__(self, in_c, out_c, kernel=(1, 1), stride=(1, 1), padding=(0, 0), groups=1):31        super(LinearBlock, self).__init__()32        self.layers = nn.Sequential(33            Conv2d(in_c, out_c, kernel, stride, padding, groups=groups, bias=False),34            BatchNorm2d(num_features=out_c)35        )36 37    def forward(self, x):38        return self.layers(x)39 40 41class DepthWise(Module):42    def __init__(self, in_c, out_c, residual=False, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=1):43        super(DepthWise, self).__init__()44        self.residual = residual45        self.layers = nn.Sequential(46            ConvBlock(in_c, out_c=groups, kernel=(1, 1), padding=(0, 0), stride=(1, 1)),47            ConvBlock(groups, groups, groups=groups, kernel=kernel, padding=padding, stride=stride),48            LinearBlock(groups, out_c, kernel=(1, 1), padding=(0, 0), stride=(1, 1))49        )50 51    def forward(self, x):52        short_cut = None53        if self.residual:54            short_cut = x55        x = self.layers(x)56        if self.residual:57            output = short_cut + x58        else:59            output = x60        return output61 62 63class Residual(Module):64    def __init__(self, c, num_block, groups, kernel=(3, 3), stride=(1, 1), padding=(1, 1)):65        super(Residual, self).__init__()66        modules = []67        for _ in range(num_block):68            modules.append(DepthWise(c, c, True, kernel, stride, padding, groups))69        self.layers = Sequential(*modules)70 71    def forward(self, x):72        return self.layers(x)73 74 75class GDC(Module):76    def __init__(self, embedding_size):77        super(GDC, self).__init__()78        self.layers = nn.Sequential(79            LinearBlock(512, 512, groups=512, kernel=(7, 7), stride=(1, 1), padding=(0, 0)),80            Flatten(),81            Linear(512, embedding_size, bias=False),82            BatchNorm1d(embedding_size))83 84    def forward(self, x):85        return self.layers(x)86 87 88class MobileFaceNet(Module):89    def __init__(self, fp16=False, num_features=512):90        super(MobileFaceNet, self).__init__()91        scale = 292        self.fp16 = fp1693        self.layers = nn.Sequential(94            ConvBlock(3, 64 * scale, kernel=(3, 3), stride=(2, 2), padding=(1, 1)),95            ConvBlock(64 * scale, 64 * scale, kernel=(3, 3), stride=(1, 1), padding=(1, 1), groups=64),96            DepthWise(64 * scale, 64 * scale, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=128),97            Residual(64 * scale, num_block=4, groups=128, kernel=(3, 3), stride=(1, 1), padding=(1, 1)),98            DepthWise(64 * scale, 128 * scale, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=256),99            Residual(128 * scale, num_block=6, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1)),100            DepthWise(128 * scale, 128 * scale, kernel=(3, 3), stride=(2, 2), padding=(1, 1), groups=512),101            Residual(128 * scale, num_block=2, groups=256, kernel=(3, 3), stride=(1, 1), padding=(1, 1)),102        )103        self.conv_sep = ConvBlock(128 * scale, 512, kernel=(1, 1), stride=(1, 1), padding=(0, 0))104        self.features = GDC(num_features)105        self._initialize_weights()106 107    def _initialize_weights(self):108        for m in self.modules():109            if isinstance(m, nn.Conv2d):110                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')111                if m.bias is not None:112                    m.bias.data.zero_()113            elif isinstance(m, nn.BatchNorm2d):114                m.weight.data.fill_(1)115                m.bias.data.zero_()116            elif isinstance(m, nn.Linear):117                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')118                if m.bias is not None:119                    m.bias.data.zero_()120 121    def forward(self, x):122        with torch.cuda.amp.autocast(self.fp16):123            x = self.layers(x)124        x = self.conv_sep(x.float() if self.fp16 else x)125        x = self.features(x)126        return x127 128 129def get_mbf(fp16, num_features):130    return MobileFaceNet(fp16, num_features)