insightface/recognition/symbol/fmnasnet.py

import sys
import os
import mxnet as mx
import mxnet.ndarray as nd
import mxnet.gluon as gluon
import mxnet.gluon.nn as nn
import mxnet.autograd as ag
import symbol_utils
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from config import config

def Act():
    if config.net_act=='prelu':
      return nn.PReLU()
    else:
      return nn.Activation(config.net_act)

def ConvBlock(channels, kernel_size, strides, **kwargs):
    out = nn.HybridSequential(**kwargs)
    with out.name_scope():
        out.add(
            nn.Conv2D(channels, kernel_size, strides=strides, padding=1, use_bias=False),
            nn.BatchNorm(scale=True),
            Act()
            #nn.Activation('relu')
        )
    return out

def Conv1x1(channels, is_linear=False, **kwargs):
    out = nn.HybridSequential(**kwargs)
    with out.name_scope():
        out.add(
            nn.Conv2D(channels, 1, padding=0, use_bias=False),
            nn.BatchNorm(scale=True)
        )
        if not is_linear:
            #out.add(nn.Activation('relu'))
            out.add(Act())
    return out

def DWise(channels, strides, kernel_size=3, **kwargs):
    out = nn.HybridSequential(**kwargs)
    with out.name_scope():
        out.add(
            nn.Conv2D(channels, kernel_size, strides=strides, padding=kernel_size // 2, groups=channels, use_bias=False),
            nn.BatchNorm(scale=True),
            Act()
            #nn.Activation('relu')
        )
    return out

class SepCONV(nn.HybridBlock):
    def __init__(self, inp, output, kernel_size, depth_multiplier=1, with_bn=True, **kwargs):
        super(SepCONV, self).__init__(**kwargs)
        with self.name_scope():
            self.net = nn.HybridSequential()
            cn = int(inp*depth_multiplier)

            if output is None:
                self.net.add(
                    nn.Conv2D(in_channels=inp, channels=cn, groups=inp, kernel_size=kernel_size, strides=(1,1), padding=kernel_size // 2
                        , use_bias=not with_bn)
                )
            else:
                self.net.add(
                    nn.Conv2D(in_channels=inp, channels=cn, groups=inp, kernel_size=kernel_size, strides=(1,1), padding=kernel_size // 2
                        , use_bias=False),
                    nn.BatchNorm(),
                    Act(),
                    #nn.Activation('relu'),
                    nn.Conv2D(in_channels=cn, channels=output, kernel_size=(1,1), strides=(1,1)
                        , use_bias=not with_bn)
                )

            self.with_bn = with_bn
            self.act = Act()
            #self.act = nn.Activation('relu')
            if with_bn:
                self.bn = nn.BatchNorm()
    def hybrid_forward(self, F ,x):
        x = self.net(x)
        if self.with_bn:
            x = self.bn(x)
        if self.act is not None:
            x = self.act(x)
        return x

class ExpandedConv(nn.HybridBlock):
    def __init__(self, inp, oup, t, strides, kernel=3, same_shape=True, **kwargs):
        super(ExpandedConv, self).__init__(**kwargs)

        self.same_shape = same_shape
        self.strides = strides
        with self.name_scope(): 
            self.bottleneck = nn.HybridSequential()
            self.bottleneck.add(
                Conv1x1(inp*t, prefix="expand_"),
                DWise(inp*t, self.strides, kernel, prefix="dwise_"),
                Conv1x1(oup, is_linear=True, prefix="linear_")
            )
    def hybrid_forward(self, F, x):
        out = self.bottleneck(x)
        if self.strides == 1 and self.same_shape:
            out = F.elemwise_add(out, x)
        return out

def ExpandedConvSequence(t, k, inp, oup, repeats, first_strides, **kwargs):
    seq = nn.HybridSequential(**kwargs)
    with seq.name_scope():
        seq.add(ExpandedConv(inp, oup, t, first_strides, k, same_shape=False))
        curr_inp = oup
        for i in range(1, repeats):
            seq.add(ExpandedConv(curr_inp, oup, t, 1))
            curr_inp = oup
    return seq

class MNasNet(nn.HybridBlock):
    def __init__(self, m=1.0, **kwargs):
        super(MNasNet, self).__init__(**kwargs)
        
        self.first_oup = int(32*m)
        self.second_oup = int(16*m)
        #self.second_oup = int(32*m)
        self.interverted_residual_setting = [
            # t, c,  n, s, k
            [3, int(24*m),  3, 2, 3, "stage2_"],  # -> 56x56
            [3, int(40*m),  3, 2, 5, "stage3_"],  # -> 28x28
            [6, int(80*m),  3, 2, 5, "stage4_1_"],  # -> 14x14
            [6, int(96*m),  2, 1, 3, "stage4_2_"],  # -> 14x14
            [6, int(192*m), 4, 2, 5, "stage5_1_"], # -> 7x7
            [6, int(320*m), 1, 1, 3, "stage5_2_"], # -> 7x7          
        ]
        self.last_channels = int(1024*m)

        with self.name_scope():
            self.features = nn.HybridSequential()
            self.features.add(ConvBlock(self.first_oup, 3, 1, prefix="stage1_conv0_"))
            self.features.add(SepCONV(self.first_oup, self.second_oup, 3, prefix="stage1_sepconv0_"))
            inp = self.second_oup
            for i, (t, c, n, s, k, prefix) in enumerate(self.interverted_residual_setting):
                oup = c
                self.features.add(ExpandedConvSequence(t, k, inp, oup, n, s, prefix=prefix))
                inp = oup

            self.features.add(Conv1x1(self.last_channels, prefix="stage5_3_"))
            #self.features.add(nn.GlobalAvgPool2D())
            #self.features.add(nn.Flatten())
            #self.output = nn.Dense(num_classes)
    def hybrid_forward(self, F, x):
        x = self.features(x)
        #x = self.output(x)
        return x

    def num_output_channel(self):
      return self.last_channels

def get_symbol():
  net = MNasNet(config.net_multiplier)
  data = mx.sym.Variable(name='data')
  data = data-127.5
  data = data*0.0078125
  body = net(data)
  fc1 = symbol_utils.get_fc1(body, config.emb_size, config.net_output, input_channel=net.num_output_channel())
  return fc1