forked from mindspore-Ecosystem/mindspore
expand bias_add and bias_add_grad op
This commit is contained in:
parent
b3b553245f
commit
febdb1850c
|
@ -18,3 +18,5 @@ from .gelu import expand_gelu
|
|||
from .layernorm import expand_layernorm
|
||||
from .softmax import expand_softmax
|
||||
from .square import expand_square
|
||||
from .bias_add import expand_biasadd
|
||||
from .bias_add_grad import expand_biasaddgrad
|
||||
|
|
|
@ -0,0 +1,62 @@
|
|||
# Copyright 2020 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.
|
||||
# ===========================================================================
|
||||
"""generate json desc for bias_add"""
|
||||
from mindspore._extends.graph_kernel.model import model_builder as builder
|
||||
|
||||
|
||||
def expand_biasadd(expand_info):
|
||||
"""BiasAdd expander"""
|
||||
|
||||
# get op info.
|
||||
input_desc_0 = expand_info['input_desc'][0]
|
||||
input_desc_1 = expand_info['input_desc'][1]
|
||||
graph_builder = builder.GraphBuilder()
|
||||
# generate a graph.
|
||||
with graph_builder.graph_scope('main') as graph_scope:
|
||||
# create tensor input.
|
||||
input_x = graph_builder.tensor(
|
||||
input_desc_0['shape'], input_desc_0['data_type'], input_desc_0['format'])
|
||||
input_y = graph_builder.tensor(
|
||||
input_desc_1['shape'], input_desc_1['data_type'], input_desc_1['format'])
|
||||
graph_scope.set_input(input_x, input_y)
|
||||
if input_x.data_format == "NCHW":
|
||||
input_y_expand = graph_builder.emit(
|
||||
'ExpandDims', [input_y], attrs={'axis': 1})
|
||||
input_y_expand = graph_builder.emit(
|
||||
'ExpandDims', [input_y_expand], attrs={'axis': 2})
|
||||
result = graph_builder.emit('TensorAdd', [input_x, input_y_expand])
|
||||
elif input_x.data_format == "DefaultFormat":
|
||||
if len(input_x.shape) == 2:
|
||||
result = graph_builder.emit('TensorAdd', [input_x, input_y])
|
||||
elif len(input_x.shape) == 3:
|
||||
input_y_expand = graph_builder.emit(
|
||||
'ExpandDims', [input_y], attrs={'axis': 1})
|
||||
result = graph_builder.emit(
|
||||
'TensorAdd', [input_x, input_y_expand])
|
||||
else:
|
||||
input_y_expand = graph_builder.emit(
|
||||
'ExpandDims', [input_y], attrs={'axis': 1})
|
||||
input_y_expand = graph_builder.emit(
|
||||
'ExpandDims', [input_y_expand], attrs={'axis': 2})
|
||||
result = graph_builder.emit(
|
||||
'TensorAdd', [input_x, input_y_expand])
|
||||
else:
|
||||
result = graph_builder.emit('TensorAdd', [input_x, input_y])
|
||||
|
||||
# set graph output.
|
||||
graph_scope.set_output(result)
|
||||
|
||||
graph = graph_builder.get()[0]
|
||||
return graph
|
|
@ -0,0 +1,48 @@
|
|||
# Copyright 2020 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.
|
||||
# ===========================================================================
|
||||
"""generate json desc for bias_add"""
|
||||
from mindspore._extends.graph_kernel.model import model_builder as builder
|
||||
|
||||
|
||||
def expand_biasaddgrad(expand_info):
|
||||
"""BiasAddGrad expander"""
|
||||
# get op info.
|
||||
input_desc_0 = expand_info['input_desc'][0]
|
||||
graph_builder = builder.GraphBuilder()
|
||||
# generate a graph.
|
||||
with graph_builder.graph_scope('main') as graph_scope:
|
||||
# create tensor input.
|
||||
input_x = graph_builder.tensor(
|
||||
input_desc_0['shape'], input_desc_0['data_type'], input_desc_0['format'])
|
||||
graph_scope.set_input(input_x)
|
||||
reduce_axis = ()
|
||||
if input_x.data_format == 'NHWC':
|
||||
reduce_axis = (0, 1, 2)
|
||||
elif input_x.data_format == 'NCHW':
|
||||
reduce_axis = (0, 2, 3)
|
||||
# Default format shape's length maybe equal 2 to 4, so different shape's length reduce axis are differnet
|
||||
else:
|
||||
if len(input_x.shape) == 2:
|
||||
reduce_axis = (0,)
|
||||
elif len(input_x.shape) == 3:
|
||||
reduce_axis = (0, 1)
|
||||
else:
|
||||
reduce_axis = (0, 2, 3)
|
||||
result = graph_builder.emit('ReduceSum', [input_x], attrs={'reduce_axis': reduce_axis, 'keep_dims': False})
|
||||
# set graph output.
|
||||
graph_scope.set_output(result)
|
||||
|
||||
graph = graph_builder.get()[0]
|
||||
return graph
|
|
@ -296,6 +296,7 @@ class Graph:
|
|||
def __init__(self, name, ops):
|
||||
self.name = name
|
||||
self.ops = ops # in topo order, can not use set
|
||||
self.inputs = []
|
||||
self.outputs = []
|
||||
|
||||
def set_processor(self, processor):
|
||||
|
@ -341,6 +342,9 @@ class Graph:
|
|||
if d not in self.ops:
|
||||
outputs.append(op.output)
|
||||
break
|
||||
if self.inputs:
|
||||
inputs = self.inputs
|
||||
|
||||
if self.outputs:
|
||||
outputs = self.outputs
|
||||
return inputs, outputs
|
||||
|
|
|
@ -22,14 +22,36 @@ class OpInfer:
|
|||
"""Op infer"""
|
||||
@staticmethod
|
||||
def default_reduce_infer(inputs, attrs):
|
||||
"""Default reduce infer"""
|
||||
shape = copy.deepcopy(inputs[0].shape)
|
||||
for i in attrs['reduce_axis']:
|
||||
shape[i] = 1
|
||||
if attrs['keep_dims']:
|
||||
for i in attrs['reduce_axis']:
|
||||
shape[i] = 1
|
||||
return shape
|
||||
|
||||
real_shape = []
|
||||
for i, _ in enumerate(shape):
|
||||
if i not in attrs['reduce_axis']:
|
||||
real_shape.append(shape[i])
|
||||
return real_shape
|
||||
|
||||
@staticmethod
|
||||
def default_elementwise_infer(inputs, attrs):
|
||||
"""Default elementwise infer"""
|
||||
shape = (1,)
|
||||
max_flatten_shape = 1
|
||||
for t in inputs:
|
||||
flatten_shape = 1
|
||||
for s in t.shape:
|
||||
flatten_shape *= s
|
||||
if flatten_shape >= max_flatten_shape:
|
||||
max_flatten_shape = flatten_shape
|
||||
shape = t.shape
|
||||
return shape
|
||||
|
||||
default_infer_shape_func = [
|
||||
None,
|
||||
lambda inputs, attrs: max([t.shape for t in inputs]),
|
||||
default_elementwise_infer.__func__,
|
||||
lambda inputs, attrs: max([t.shape for t in inputs]),
|
||||
default_reduce_infer.__func__,
|
||||
None,
|
||||
|
@ -72,9 +94,16 @@ class OpInfer:
|
|||
class GraphBuilder:
|
||||
"""Graph builder"""
|
||||
class GraphWrapper:
|
||||
"""Graph wrapper"""
|
||||
|
||||
def __init__(self, name):
|
||||
self.graph = Graph(name, [])
|
||||
|
||||
def set_input(self, *para):
|
||||
for t in para:
|
||||
t.para_type = Tensor.PARA_INPUT
|
||||
self.graph.inputs.append(t)
|
||||
|
||||
def set_output(self, *para):
|
||||
for t in para:
|
||||
t.para_type = Tensor.PARA_OUTPUT
|
||||
|
|
|
@ -702,6 +702,8 @@ FuncGraphPtr JsonDescToAnf(const std::string &json_desc, const std::vector<AnfNo
|
|||
std::unordered_set<PrimitivePtr> GetExpandOps() {
|
||||
std::unordered_set<PrimitivePtr> expand_ops = {
|
||||
prim::kPrimSquare,
|
||||
prim::kPrimBiasAdd,
|
||||
prim::kPrimBiasAddGrad,
|
||||
};
|
||||
return expand_ops;
|
||||
}
|
||||
|
|
|
@ -143,6 +143,7 @@ inline const PrimitivePtr kPrimDepthwiseConv2dNativeBackpropFilter =
|
|||
std::make_shared<Primitive>("DepthwiseConv2dNativeBackpropFilter");
|
||||
inline const PrimitivePtr kPrimDepthwiseConv2dNativeBackpropInput =
|
||||
std::make_shared<Primitive>("DepthwiseConv2dNativeBackpropInput");
|
||||
inline const PrimitivePtr kPrimBiasAdd = std::make_shared<Primitive>("BiasAdd");
|
||||
inline const PrimitivePtr kPrimBiasAddGrad = std::make_shared<Primitive>("BiasAddGrad");
|
||||
inline const PrimitivePtr kPrimSoftmaxCrossEntropyWithLogits =
|
||||
std::make_shared<Primitive>("SoftmaxCrossEntropyWithLogits");
|
||||
|
|
Loading…
Reference in New Issue