Dubhe/dubhe_data_process/of_model/resnet_model.py

"""
/**
* Copyright 2020 Tianshu AI Platform. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* =============================================================
*/
  Reference:
    - [Identity Mappings in Deep Residual Networks]
      (https://arxiv.org/abs/1603.05027) (CVPR 2016)
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import oneflow as flow

BLOCK_COUNTS = [3, 4, 6, 3]
BLOCK_FILTERS = [256, 512, 1024, 2048]
BLOCK_FILTERS_INNER = [64, 128, 256, 512]


class ResnetBuilder(object):
    def __init__(self, weight_regularizer, trainable=True, training=True):
        self.weight_initializer = flow.variance_scaling_initializer(
            2, 'fan_in', 'random_normal', data_format="NCHW")
        self.weight_regularizer = weight_regularizer
        self.trainable = trainable
        self.training = training

    def _conv2d(
        self,
        name,
        input,
        filters,
        kernel_size,
        strides=1,
        padding="SAME",
        data_format="NCHW",
        dilations=1,
    ):
        weight = flow.get_variable(
            name + "-weight",
            shape=(filters, input.shape[1], kernel_size, kernel_size),
            dtype=input.dtype,
            initializer=self.weight_initializer,
            regularizer=self.weight_regularizer,
            model_name="weight",
            trainable=self.trainable,
        )
        return flow.nn.conv2d(
            input,
            weight,
            strides,
            padding,
            data_format,
            dilations,
            name=name)

    def _batch_norm(self, inputs, name=None, last=False):
        initializer = flow.zeros_initializer() if last else flow.ones_initializer()
        return flow.layers.batch_normalization(
            inputs=inputs,
            axis=1,
            momentum=0.9,  # 97,
            epsilon=1e-5,
            center=True,
            scale=True,
            trainable=self.trainable,
            training=self.training,
            gamma_initializer=initializer,
            moving_variance_initializer=initializer,
            gamma_regularizer=self.weight_regularizer,
            beta_regularizer=self.weight_regularizer,
            name=name,
        )

    def conv2d_affine(
            self,
            input,
            name,
            filters,
            kernel_size,
            strides,
            activation=None,
            last=False):
        # input data_format must be NCHW, cannot check now
        padding = "SAME" if strides > 1 or kernel_size > 1 else "VALID"
        output = self._conv2d(
            name,
            input,
            filters,
            kernel_size,
            strides,
            padding)
        output = self._batch_norm(output, name + "_bn", last=last)
        if activation == "Relu":
            output = flow.nn.relu(output)

        return output

    def bottleneck_transformation(
            self,
            input,
            block_name,
            filters,
            filters_inner,
            strides):
        a = self.conv2d_affine(
            input,
            block_name +
            "_branch2a",
            filters_inner,
            1,
            1,
            activation="Relu")

        b = self.conv2d_affine(
            a,
            block_name +
            "_branch2b",
            filters_inner,
            3,
            strides,
            activation="Relu")

        c = self.conv2d_affine(
            b,
            block_name +
            "_branch2c",
            filters,
            1,
            1,
            last=True)

        return c

    def residual_block(
            self,
            input,
            block_name,
            filters,
            filters_inner,
            strides_init):
        if strides_init != 1 or block_name == "res2_0":
            shortcut = self.conv2d_affine(
                input, block_name + "_branch1", filters, 1, strides_init
            )
        else:
            shortcut = input

        bottleneck = self.bottleneck_transformation(
            input, block_name, filters, filters_inner, strides_init,
        )

        return flow.nn.relu(bottleneck + shortcut)

    def residual_stage(
            self,
            input,
            stage_name,
            counts,
            filters,
            filters_inner,
            stride_init=2):
        output = input
        for i in range(counts):
            block_name = "%s_%d" % (stage_name, i)
            output = self.residual_block(
                output,
                block_name,
                filters,
                filters_inner,
                stride_init if i == 0 else 1)

        return output

    def resnet_conv_x_body(self, input):
        output = input
        for i, (counts, filters, filters_inner) in enumerate(
            zip(BLOCK_COUNTS, BLOCK_FILTERS, BLOCK_FILTERS_INNER)
        ):
            stage_name = "res%d" % (i + 2)
            output = self.residual_stage(
                output,
                stage_name,
                counts,
                filters,
                filters_inner,
                1 if i == 0 else 2)

        return output

    def resnet_stem(self, input):
        conv1 = self._conv2d("conv1", input, 64, 7, 2)
        conv1_bn = flow.nn.relu(self._batch_norm(conv1, "conv1_bn"))
        pool1 = flow.nn.max_pool2d(
            conv1_bn,
            ksize=3,
            strides=2,
            padding="SAME",
            data_format="NCHW",
            name="pool1",
        )
        return pool1


def resnet50(
        images,
        trainable=True,
        need_transpose=False,
        training=True,
        wd=1.0 / 32768):
    weight_regularizer = flow.regularizers.l2(
        wd) if wd > 0.0 and wd < 1.0 else None
    builder = ResnetBuilder(weight_regularizer, trainable, training)
    # note: images.shape = (N C H W) in cc's new dataloader, transpose is not
    # needed anymore
    if need_transpose:
        images = flow.transpose(images, name="transpose", perm=[0, 3, 1, 2])

    with flow.deprecated.variable_scope("Resnet"):
        stem = builder.resnet_stem(images)
        body = builder.resnet_conv_x_body(stem)
        pool5 = flow.nn.avg_pool2d(
            body,
            ksize=7,
            strides=1,
            padding="VALID",
            data_format="NCHW",
            name="pool5",
        )
        fc1001 = flow.layers.dense(
            flow.reshape(pool5, (pool5.shape[0], -1)),
            units=1000,
            use_bias=True,
            kernel_initializer=flow.variance_scaling_initializer(2, 'fan_in', 'random_normal'),
            bias_initializer=flow.zeros_initializer(),
            kernel_regularizer=weight_regularizer,
            bias_regularizer=weight_regularizer,
            trainable=trainable,
            name="fc1001",
        )

    return fc1001