# 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. from __future__ import annotations from typing import TYPE_CHECKING import numpy as np import paddle from paddle import base from paddle.base.layer_helper import LayerHelper from paddle.base.param_attr import ParamAttr from paddle.nn import ( Layer, initializer as I, ) if TYPE_CHECKING: from paddle import Tensor from paddle._typing import DataLayout2D, ParamAttrLike def resnet_unit( x: Tensor, filter_x: Tensor, scale_x: Tensor, bias_x: Tensor, mean_x: Tensor, var_x: Tensor, z: Tensor | None, filter_z: Tensor | None, scale_z: Tensor | None, bias_z: Tensor | None, mean_z: Tensor | None, var_z: Tensor | None, stride: int, stride_z: int, padding: int, dilation: int, groups: int, momentum: float, eps: float, data_format: DataLayout2D, fuse_add: bool, has_shortcut: bool, use_global_stats: bool, is_test: bool, act: str, ) -> Tensor: helper = LayerHelper('resnet_unit', **locals()) bn_param_dtype = base.core.VarDesc.VarType.FP32 bit_mask_dtype = base.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 performance by using cudnnv8 API. """ def __init__( self, num_channels_x: int, num_filters: int, filter_size: int, stride: int = 1, momentum: float = 0.9, eps: float = 1e-5, data_format: DataLayout2D = 'NHWC', act: str = 'relu', fuse_add: bool = False, has_shortcut: bool = False, use_global_stats: bool = False, is_test: bool = False, filter_x_attr: ParamAttrLike | None = None, scale_x_attr: ParamAttrLike | None = None, bias_x_attr: ParamAttrLike | None = None, moving_mean_x_name: str | None = None, moving_var_x_name: str | None = None, num_channels_z: int = 1, stride_z: int = 1, filter_z_attr: ParamAttrLike | None = None, scale_z_attr: ParamAttrLike | None = None, bias_z_attr: ParamAttrLike | None = None, moving_mean_z_name: str | None = None, moving_var_z_name: str | None = None, ) -> None: super().__init__() self._stride = stride self._stride_z = stride_z self._dilation = 1 self._kernel_size = paddle.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( f"conv_format must be one of {valid_format}, but got conv_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 = base.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: Tensor, z: Tensor | None = None) -> Tensor: 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