googlenet_model.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File: googlenet_model.py

import tensorflow as tf
from tensorflow.contrib.layers import variance_scaling_initializer
from tensorpack.models import *
from tensorpack.tfutils.argscope import argscope, get_arg_scope

from learned_quantization import Conv2DQuant, getBNReLUQuant, QuantizedActiv


def inception_block(l, name, ch_1x1, ch_3x3, ch_5x5, is_last_block=False, is_last=False):
    data_format = get_arg_scope()['Conv2DQuant']['data_format']
    with tf.variable_scope(name):
        conv1x1 = Conv2DQuant('1x1', l, ch_1x1, 1, nl=getBNReLUQuant if not is_last_block else tf.identity)
        conv3x3_reduce = Conv2DQuant('3x3_reduce', l, ch_3x3, 1, nl=getBNReLUQuant)
        conv3x3 = Conv2DQuant('3x3', conv3x3_reduce, ch_3x3, 3, nl=getBNReLUQuant if not is_last_block else tf.identity)
        conv5x5_reduce = Conv2DQuant('5x5_reduce', l, ch_5x5, 1, nl=getBNReLUQuant)
        conv5x5 = Conv2DQuant('5x5', conv5x5_reduce, ch_5x5, 5, nl=getBNReLUQuant if not is_last_block else tf.identity)
        if is_last_block and not is_last:
            conv1x1 = MaxPooling('pool_1x1', conv1x1, shape=3, stride=2, padding='SAME')
            conv1x1 = BNReLU('conv1x1_bn', conv1x1)
            conv1x1 = QuantizedActiv('conv1x1_quant', conv1x1)
            conv3x3 = MaxPooling('pool_3x3', conv3x3, shape=3, stride=2, padding='SAME')
            conv3x3 = BNReLU('conv3x3_bn', conv3x3)
            conv3x3 = QuantizedActiv('conv3x3_quant', conv3x3)
            conv5x5 = MaxPooling('pool_5x5', conv5x5, shape=3, stride=2, padding='SAME')
            conv5x5 = BNReLU('conv5x5_bn', conv5x5)
            conv5x5 = QuantizedActiv('conv5x5_quant', conv5x5)
        l = tf.concat([
            conv1x1,
            conv3x3,
            conv5x5], 1 if data_format == 'NCHW' else 3, name='concat')
        if is_last:
            l = BNReLU('output_bn', l)
    return l


def googlenet_backbone(image, qw=1):
    with argscope(Conv2DQuant, nl=tf.identity, use_bias=False,
                  W_init=variance_scaling_initializer(mode='FAN_IN'),
                  data_format=get_arg_scope()['Conv2D']['data_format'],
                  nbit=qw,
                  is_quant=True if qw > 0 else False):
        logits = (LinearWrap(image)
                  .Conv2DQuant('conv1', 64, 7, stride=2, is_quant=False)
                  .MaxPooling('pool1', shape=3, stride=2, padding='SAME')
                  .BNReLUQuant('pool1/out')
                  .Conv2DQuant('conv2/3x3_reduce', 192, 1, nl=getBNReLUQuant)
                  .Conv2DQuant('conv2/3x3', 192, 3)
                  .MaxPooling('pool2', shape=3, stride=2, padding='SAME')
                  .BNReLUQuant('pool2/out')
                  .apply(inception_block, 'incpetion_3a', 96, 128, 32)
                  .apply(inception_block, 'incpetion_3b', 192, 192, 96, is_last_block=True)
                  .apply(inception_block, 'incpetion_4a', 256, 208, 48)
                  .apply(inception_block, 'incpetion_4b', 224, 224, 64)
                  .apply(inception_block, 'incpetion_4c', 192, 256, 64)
                  .apply(inception_block, 'incpetion_4d', 176, 288, 64)
                  .apply(inception_block, 'incpetion_4e', 384, 320, 128, is_last_block=True)
                  .apply(inception_block, 'incpetion_5a', 384, 320, 128)
                  .apply(inception_block, 'incpetion_5b', 512, 384, 128, is_last_block=True, is_last=True)
                  .GlobalAvgPooling('pool5')
                  .FullyConnected('linear', out_dim=1000, nl=tf.identity)())
    return logits