resnet_unit.py

#   Copyright (c) 2021 PaddlePaddle Authors. 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.

import numpy as np
import paddle.fluid as fluid
from paddle.nn import initializer as I
from paddle.nn import Layer
from paddle.fluid.layers import utils
from paddle.fluid.layer_helper import LayerHelper
from paddle.fluid.param_attr import ParamAttr


def resnet_unit(x, filter_x, scale_x, bias_x, mean_x, var_x, z, filter_z,
                scale_z, bias_z, mean_z, var_z, stride, stride_z, padding,
                dilation, groups, momentum, eps, data_format, fuse_add,
                has_shortcut, use_global_stats, is_test, act):

    helper = LayerHelper('resnet_unit', **locals())
    bn_param_dtype = fluid.core.VarDesc.VarType.FP32
    bit_mask_dtype = fluid.core.VarDesc.VarType.INT32
    out = helper.create_variable_for_type_inference(x.dtype)
    bit_mask = helper.create_variable_for_type_inference(dtype=bit_mask_dtype,
                                                         stop_gradient=True)
    # intermediate_out for x
    conv_x = helper.create_variable_for_type_inference(dtype=x.dtype,
                                                       stop_gradient=True)
    saved_mean_x = helper.create_variable_for_type_inference(
        dtype=bn_param_dtype, stop_gradient=True)
    saved_invstd_x = helper.create_variable_for_type_inference(
        dtype=bn_param_dtype, stop_gradient=True)
    running_mean_x = mean_x
    running_var_x = var_x
    # intermediate_out for z
    conv_z = helper.create_variable_for_type_inference(dtype=x.dtype,
                                                       stop_gradient=True)
    saved_mean_z = helper.create_variable_for_type_inference(
        dtype=bn_param_dtype, stop_gradient=True)
    saved_invstd_z = helper.create_variable_for_type_inference(
        dtype=bn_param_dtype, stop_gradient=True)
    running_mean_z = helper.create_variable_for_type_inference(
        dtype=bn_param_dtype, stop_gradient=True) if mean_z is None else mean_z
    running_var_z = helper.create_variable_for_type_inference(
        dtype=bn_param_dtype, stop_gradient=True) if var_z is None else var_z

    inputs = {
        'X': x,
        'FilterX': filter_x,
        'ScaleX': scale_x,
        'BiasX': bias_x,
        'MeanX': mean_x,
        'VarX': var_x,
        'Z': z,
        'FilterZ': filter_z,
        'ScaleZ': scale_z,
        'BiasZ': bias_z,
        'MeanZ': mean_z,
        'VarZ': var_z
    }

    attrs = {
        'stride': stride,
        'stride_z': stride_z,
        'padding': padding,
        'dilation': dilation,
        'group': groups,
        'momentum': momentum,
        'epsilon': eps,
        'data_format': data_format,
        'fuse_add': fuse_add,
        'has_shortcut': has_shortcut,
        'use_global_stats': use_global_stats,
        'is_test': is_test,
        'act_type': act
    }

    outputs = {
        'Y': out,
        'BitMask': bit_mask,
        'ConvX': conv_x,
        'SavedMeanX': saved_mean_x,
        'SavedInvstdX': saved_invstd_x,
        'RunningMeanX': running_mean_x,
        'RunningVarX': running_var_x,
        'ConvZ': conv_z,
        'SavedMeanZ': saved_mean_z,
        'SavedInvstdZ': saved_invstd_z,
        'RunningMeanZ': running_mean_z,
        'RunningVarZ': running_var_z,
    }

    helper.append_op(type='resnet_unit',
                     inputs=inputs,
                     outputs=outputs,
                     attrs=attrs)

    return out


class ResNetUnit(Layer):
    r"""
    ******Temporary version******.
    ResNetUnit is designed for optimize the performence by using cudnnv8 API.
    """

    def __init__(self,
                 num_channels_x,
                 num_filters,
                 filter_size,
                 stride=1,
                 momentum=0.9,
                 eps=1e-5,
                 data_format='NHWC',
                 act='relu',
                 fuse_add=False,
                 has_shortcut=False,
                 use_global_stats=False,
                 is_test=False,
                 filter_x_attr=None,
                 scale_x_attr=None,
                 bias_x_attr=None,
                 moving_mean_x_name=None,
                 moving_var_x_name=None,
                 num_channels_z=1,
                 stride_z=1,
                 filter_z_attr=None,
                 scale_z_attr=None,
                 bias_z_attr=None,
                 moving_mean_z_name=None,
                 moving_var_z_name=None):
        super(ResNetUnit, self).__init__()
        self._stride = stride
        self._stride_z = stride_z
        self._dilation = 1
        self._kernel_size = utils.convert_to_list(filter_size, 2, 'kernel_size')
        self._padding = (filter_size - 1) // 2
        self._groups = 1
        self._momentum = momentum
        self._eps = eps
        self._data_format = data_format
        self._act = act
        self._fuse_add = fuse_add
        self._has_shortcut = has_shortcut
        self._use_global_stats = use_global_stats
        self._is_test = is_test

        # check format
        valid_format = {'NHWC', 'NCHW'}
        if data_format not in valid_format:
            raise ValueError(
                "conv_format must be one of {}, but got conv_format='{}'".
                format(valid_format, data_format))

        def _get_default_param_initializer(channels):
            filter_elem_num = np.prod(self._kernel_size) * channels
            std = (2.0 / filter_elem_num)**0.5
            return I.Normal(0.0, std)

        is_nchw = (data_format == 'NCHW')
        # initial filter
        bn_param_dtype = fluid.core.VarDesc.VarType.FP32
        if not is_nchw:
            bn_param_shape = [1, 1, 1, num_filters]
            filter_x_shape = [
                num_filters, filter_size, filter_size, num_channels_x
            ]
            filter_z_shape = [
                num_filters, filter_size, filter_size, num_channels_z
            ]
        else:
            bn_param_shape = [1, num_filters, 1, 1]
            filter_x_shape = [
                num_filters, num_channels_x, filter_size, filter_size
            ]
            filter_z_shape = [
                num_filters, num_channels_z, filter_size, filter_size
            ]

        self.filter_x = self.create_parameter(
            shape=filter_x_shape,
            attr=filter_x_attr,
            default_initializer=_get_default_param_initializer(num_channels_x))
        self.scale_x = self.create_parameter(
            shape=bn_param_shape,
            attr=scale_x_attr,
            dtype=bn_param_dtype,
            default_initializer=I.Constant(1.0))
        self.bias_x = self.create_parameter(shape=bn_param_shape,
                                            attr=bias_x_attr,
                                            dtype=bn_param_dtype,
                                            is_bias=True)
        self.mean_x = self.create_parameter(attr=ParamAttr(
            name=moving_mean_x_name,
            initializer=I.Constant(0.0),
            trainable=False),
                                            shape=bn_param_shape,
                                            dtype=bn_param_dtype)
        self.mean_x.stop_gradient = True
        self.var_x = self.create_parameter(attr=ParamAttr(
            name=moving_var_x_name,
            initializer=I.Constant(1.0),
            trainable=False),
                                           shape=bn_param_shape,
                                           dtype=bn_param_dtype)
        self.var_x.stop_gradient = True
        if has_shortcut:
            self.filter_z = self.create_parameter(
                shape=filter_z_shape,
                attr=filter_z_attr,
                default_initializer=_get_default_param_initializer(
                    num_channels_z))
            self.scale_z = self.create_parameter(
                shape=bn_param_shape,
                attr=scale_z_attr,
                dtype=bn_param_dtype,
                default_initializer=I.Constant(1.0))
            self.bias_z = self.create_parameter(shape=bn_param_shape,
                                                attr=bias_z_attr,
                                                dtype=bn_param_dtype,
                                                is_bias=True)
            self.mean_z = self.create_parameter(attr=ParamAttr(
                name=moving_mean_z_name,
                initializer=I.Constant(0.0),
                trainable=False),
                                                shape=bn_param_shape,
                                                dtype=bn_param_dtype)
            self.mean_z.stop_gradient = True
            self.var_z = self.create_parameter(attr=ParamAttr(
                name=moving_var_z_name,
                initializer=I.Constant(1.0),
                trainable=False),
                                               shape=bn_param_shape,
                                               dtype=bn_param_dtype)
            self.var_z.stop_gradient = True
        else:
            self.filter_z = None
            self.scale_z = None
            self.bias_z = None
            self.mean_z = None
            self.var_z = None

    def forward(self, x, z=None):
        if self._fuse_add and z is None:
            raise ValueError("z can not be None")

        out = resnet_unit(x, self.filter_x, self.scale_x, self.bias_x,
                          self.mean_x, self.var_x, z, self.filter_z,
                          self.scale_z, self.bias_z, self.mean_z, self.var_z,
                          self._stride, self._stride_z, self._padding,
                          self._dilation, self._groups, self._momentum,
                          self._eps, self._data_format, self._fuse_add,
                          self._has_shortcut, self._use_global_stats,
                          self._is_test, self._act)
        return out