|
- #!/usr/bin/env python3
- # coding: utf-8
- # Copyright 2019 Huawei Technologies Co., Ltd
- #
- # 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.
-
- """operator dsl function: conv_input_ad"""
- import akg.tvm
- import akg.topi
- import akg
- from akg.ops.nn import conv_backprop_input
- from akg.ops.nn import conv as conv_forward
- from akg.utils.format_transform import tvm_array_to_list
- from akg.utils import validation_check as vc_util
-
-
- def expr_to_int(in_expr):
- """Converte expr to int type value."""
- result = [a.value for a in in_expr]
- return result
-
-
- @akg.tvm.register_func("akg.autodiff.conv_input_ad_tensor")
- def conv_input_ad_tensor(data, fmap_shape, filter_shape, pad_, stride_, dilation_, attrs=None):
- """wraper of convolution filter backprop func."""
- data_list = tvm_array_to_list(data)
- fmap_shape = expr_to_int(fmap_shape)
- filter_shape = expr_to_int(filter_shape)
- pad_ = expr_to_int(pad_)
- stride_ = expr_to_int(stride_)
- dilation_ = expr_to_int(dilation_)
- c, _ = conv_backprop_input.conv_backprop_input(data_list, fmap_shape, filter_shape,
- pad_, stride_, dilation_, attrs=attrs)
- return c
-
-
- def conv_input_ad_config(data, fmap_shape, filter_shape, pad_, stride_, dilation_, attrs=None):
- """Configuration of convolution filter gradient."""
- _, configs = conv_backprop_input.conv_backprop_input(data, fmap_shape, filter_shape,
- pad_, stride_, dilation_, attrs=attrs)
- return configs
-
-
- @vc_util.check_input_type((list, tuple), (list, tuple), (list, tuple),
- (list, tuple), (list, tuple), (list, tuple), (dict, type(None)))
- def conv_input_ad(input_ad_inputs, fmap_shape, filter_shape, pad_, stride_, dilation_, attrs=None):
- """
- Compute dx according to "conv forward".
-
- Args:
- input_ad_inputs (list[tvm.tensor.Tensor]): a list with length 2.
- input_ad_inputs[0](consider as dy) Tensor of type float16 ,shape 5D(out_n, out_c//C0, out_h, out_w,C0)
- input_ad_inputs[1](consider as w) Tensor of type float16 ,shape 4D(wC//C0*wH*wW, wN//C0, C0,C0)
- fmap_shape (list): [fN, fC, fH, fW]
- filter_shape (list): [wN, wC, wH, wW]
- pad_ (list): [pad_left, pad_right, pad_top, pad_bottom]
- stride_ (list): [stride_h, stride_w]
- dilation_ (list): [dilation_h, dilation_w]
- attrs (dict): a dict with keys like conv_tile, bypass and etc.
-
- Returns:
- tvm.tensor.Tensor, configs.
- """
-
- backward_dy, forward_w = input_ad_inputs
-
- in_n, in_c, in_h, in_w = fmap_shape
- block_size = 16
- in_c = (in_c + block_size - 1) // block_size * block_size
- x_5d_shape = (in_n, in_c // block_size, in_h, in_w, block_size)
-
- forward_x = akg.tvm.placeholder(x_5d_shape, forward_w.dtype, "input_X")
- original_filter_shape = akg.tvm.placeholder(filter_shape, forward_w.dtype, "input_filter")
- forward_output, _ = conv_forward.conv([forward_x, forward_w], fmap_shape, filter_shape,
- pad_, stride_, dilation_, use_bias=False, attrs=attrs)
-
- ad_attrs = {"ad_conv_enable": 1, "ad_conv_reuse_conv": 0}
- jacs = list(akg.differentiate(forward_output, [forward_x], backward_dy, ad_attrs,
- [backward_dy, forward_w, original_filter_shape]))
- configs = conv_input_ad_config([backward_dy, forward_w], fmap_shape, filter_shape,
- pad_, stride_, dilation_, attrs=attrs)
-
- return jacs[0], configs
|
