EfficientNet

gistfile1.txt

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import math
import random
import chainer
import chainer.functions as F
import chainer.links as L
import numpy as np
from chainer import reporter
initializer = chainer.initializers.HeNormal()

bn_decay=0.99
bn_eps=1e-3

class BottleNeckA(chainer.Chain):
    def __init__(self, k=3, s=1, e=1, i=32, o=16, se=0.25):
        w = math.sqrt(2)
        self.e=e
        if e!=1:
            super(BottleNeckA, self).__init__(
                conv1=L.Convolution2D(i,  i*e, 1, 1, 0, nobias=True, initialW=initializer),
                bn1=L.BatchNormalization(i*e, decay=bn_decay, eps=bn_eps),
                conv2=L.Convolution2D(i*e, i*e, k, s, int(k/2), nobias=True, initialW=initializer, groups=i*e),
                bn2=L.BatchNormalization(i*e, decay=bn_decay, eps=bn_eps),
                se1=L.Convolution2D(i*e, int(i*se), 1, 1, 0, initialW=initializer),
                se2=L.Convolution2D(int(i*se), i*e, 1, 1, 0, initialW=initializer),
                conv3=L.Convolution2D(i*e, o, 1, 1, 0, nobias=True, initialW=initializer),
                bn3=L.BatchNormalization(o, decay=bn_decay, eps=bn_eps),
            )
        else:
            super(BottleNeckA, self).__init__(
                conv2=L.Convolution2D(i*e, i*e, k, s, int(k/2), nobias=True, initialW=initializer, groups=i*e),
                bn2=L.BatchNormalization(i*e, decay=bn_decay, eps=bn_eps),
                se1=L.Convolution2D(i*e, int(i*se), 1, 1, 0, initialW=initializer),
                se2=L.Convolution2D(int(i*se), i*e, 1, 1, 0, initialW=initializer),
                conv3=L.Convolution2D(i*e, o, 1, 1, 0, nobias=True, initialW=initializer),
                bn3=L.BatchNormalization(o, decay=bn_decay, eps=bn_eps),
            )

    def __call__(self, x):
        if self.e!=1:
            h = self.conv1(x)
            h = self.bn1(h)
            h = h*F.sigmoid(h)#swish
        else:
            h = x
        
        h = self.conv2(h)
        h = self.bn2(h)
        h = h*F.sigmoid(h)#swish
        
        se = F.mean(h, axis=(2,3), keepdims=True)
        se = self.se1(se)
        se = se*F.sigmoid(se)#swish
        se = self.se2(se)
        se = F.sigmoid(se)
        se = F.broadcast_to(se,h.shape)
        h = h*se
        
        h = self.conv3(h)
        h = self.bn3(h)
        
        if x.shape[1]==h.shape[1]:
            h=x+h
        return h

class MBConv(chainer.Chain):
    def __init__(self, r=1, k=3, s=1, e=1, i=32, o=16, se=0.25):
        #r: repeat number of MBConv stage
        #k: kernel size of DepthWiseConv
        #s: stride size of DepthWiseConv
        #e: expand ratio for invertive bottkeneck
        #i: input channel
        #o: output channel
        #se: Sqeeze Excitation ratio
        super(MBConv, self).__init__()
        links = [('a', BottleNeckA(k=k, s=s, e=e, i=i, o=o, se=se))]
        for i in range(r-1):
            links += [('b{}'.format(i+1), BottleNeckA(k=k, s=1, e=e, i=o, o=o, se=se))]

        for link in links:
            self.add_link(*link)
        self.forward = links

    def __call__(self, x):
        for name,_ in self.forward:
            f = getattr(self, name)
            h = f(x if name == 'a' else h)
        return h

class EfficientNet_B0(chainer.Chain):

    insize = 224

    def __init__(self, n_class):
        self.dropout=0.2
        super(EfficientNet_B0, self).__init__(
            conv1=L.Convolution2D(3, 32, 3, 2, 1, nobias=True, initialW=initializer),
            bn1=L.BatchNormalization(32, decay=bn_decay, eps=bn_eps),
            MB2=MBConv(r=1,k=3,s=1,e=1,i=32,o=16,se=0.25),
            MB3=MBConv(r=2,k=3,s=2,e=6,i=16,o=24,se=0.25),
            MB4=MBConv(r=2,k=5,s=2,e=6,i=24,o=40,se=0.25),
            MB5=MBConv(r=3,k=3,s=2,e=6,i=40,o=80,se=0.25),
            MB6=MBConv(r=3,k=5,s=1,e=6,i=80,o=112,se=0.25),
            MB7=MBConv(r=4,k=5,s=2,e=6,i=112,o=192,se=0.25),
            MB8=MBConv(r=1,k=3,s=1,e=6,i=192,o=320,se=0.25),
            conv9=L.Convolution2D(None, 1280, 1, 1, 0, nobias=True, initialW=initializer),
            bn9=L.BatchNormalization(1280, decay=bn_decay, eps=bn_eps),
            fc=L.Linear(None, n_class, initialW=initializer),
        )

    def __call__(self, x, t):
        h = self.conv1(x)
        h = self.bn1(h)
        h = h*F.sigmoid(h)#swish
        
        h = self.MB2(h)
        h = self.MB3(h)
        h = self.MB4(h)
        h = self.MB5(h)
        h = self.MB6(h)
        h = self.MB7(h)
        h = self.MB8(h)
        
        h = self.bn9(self.conv9(x))
        h = h*F.sigmoid(h)#swish
        
        h = F.mean(h, axis=(2,3))
        if chainer.config.train:
            if self.dropout is not None:
                h=F.dropout(h,self.dropout)
        h = self.fc(h)
        
        if chainer.config.train:
            self.loss = (F.sum(t[t.nonzero()] * F.log(t[t.nonzero()])) - F.sum(t * F.log_softmax(h))) / h.shape[0]
            self.accuracy = chainer.functions.evaluation.accuracy.accuracy(h,F.argmax(t,1))
            reporter.report({'loss': self.loss}, self)
            reporter.report({'accuracy': self.accuracy}, self)
        else:
            self.loss = chainer.functions.loss.softmax_cross_entropy.softmax_cross_entropy(h,t)
            self.accuracy = chainer.functions.evaluation.accuracy.accuracy(h,t)
            reporter.report({'loss': self.loss}, self)
            reporter.report({'accuracy': self.accuracy}, self)
        
        return self.loss