Skip to content

onnx frontend support layout choice depend on hardware target supported layout with NCHW and NHWC #4474

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Beya2019 wants to merge 2 commits into from
Closed

onnx frontend support layout choice depend on hardware target supported layout with NCHW and NHWC #4474

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c23e15e
Merge pull request #1 from apache/master
Beya2019 Dec 5, 2019
7f5eccc
onnx frontend support layout choice depend on hardware target support…
Dec 6, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions python/tvm/relay/frontend/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -393,9 +393,6 @@ def __call__(self, inputs, attrs, *args):
' supported.', k, op_name)
elif k in self._disables:
logging.warning("Attribute %s is disabled in relay.sym.%s", k, op_name)
elif k in self._ignores:
if k != 'tvm_custom':
logging.warning("Attribute %s is ignored in relay.sym.%s", k, op_name)
elif k in self._transforms:
new_name, defaults, transform = self._parse_default(self._transforms[k])
if defaults is None:
Expand All @@ -406,6 +403,9 @@ def __call__(self, inputs, attrs, *args):
new_attrs[new_name] = defaults
else:
new_attrs[new_name] = transform(new_attr)
elif k in self._ignores:
if k != 'tvm_custom':
logging.warning("Attribute %s is ignored in relay.sym.%s", k, op_name)
else:
# copy
new_attrs[k] = attrs[k]
Expand Down Expand Up @@ -545,4 +545,8 @@ def __init__(self, new_name):
def __call__(self, inputs, attrs, *args):
if 'tvm_custom' in attrs:
attrs.pop('tvm_custom')
if "global_" in self._new_name:
attrs['layout'] = attrs['_target_layout']
if '_target_layout' in attrs:
attrs.pop('_target_layout')
return get_relay_op(self._new_name)(*inputs, **attrs)
56 changes: 49 additions & 7 deletions python/tvm/relay/frontend/onnx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,13 @@ def _dim_check(attrs):

return _dim_check, "Only 2d kernel supported."

def transform_layout_axis(layout, axis):
"""onnx transform axis depend the layout 'NCHW' and 'NHWC'."""
if layout == 'NHWC' and axis == 1:
return 3
else:
return axis


class OnnxOpConverter(object):
""" A helper class for holding onnx op converters.
Expand Down Expand Up @@ -127,7 +134,7 @@ def _impl_v1(cls, inputs, attr, params):
conv_ops = ["conv2d", "conv2d_transpose"]
if attr.get('broadcast', 0) and any(x in str(inputs[0]) for x in conv_ops):
# TODO(zhreshold): remove hard coded infershape
axis = int(attr.get('axis', 0))
axis = 0 if attr['_target_layout'] == 'NHWC' else int(attr.get('axis', 0))
inputs[1] = _op.expand_dims(inputs[1], axis=axis, num_newaxis=2)
return get_relay_op(op_name)(*inputs)

Expand All @@ -143,6 +150,7 @@ def _impl_v1(cls, inputs, attr, params):
op_name=dimension_picker(cls.name),
transforms={
'kernel_shape': 'pool_size',
'_target_layout': 'layout',
'pads': ('padding', (0, 0), revert_caffe2_pad)
},
# very weird attributes here in onnx, force check
Expand Down Expand Up @@ -180,6 +188,8 @@ class BatchNorm(OnnxOpConverter):
@classmethod
def _impl_v1(cls, inputs, attr, params):
# TODO(zhreshold): 'spatial' is not properly handled here.
axis = attr.get('axis', 1)
attr['axis'] = transform_layout_axis(attr['_target_layout'], axis)
out = AttrCvt(
op_name='batch_norm',
ignores=['spatial', 'is_test', 'consumed_inputs', 'momentum'])(inputs, attr,
Expand All @@ -202,17 +212,27 @@ class Conv(OnnxOpConverter):

@classmethod
def _impl_v1(cls, inputs, attr, params):
if attr['_target_layout'] == 'NHWC':
if attr['group'] != 1:
#depthwise conv2d
attr['kernel_layout'] = 'HWOI'
inputs[1] = _op.transpose(inputs[1], axes=(2, 3, 0, 1))
else:
attr['kernel_layout'] = 'HWIO'
inputs[1] = _op.transpose(inputs[1], axes=(2, 3, 1, 0))
out = AttrCvt(op_name=dimension_picker('conv'),
transforms={
'kernel_shape': 'kernel_size',
'dilations': ('dilation', (0, 0)),
'pads': ('padding', (0, 0), revert_caffe2_pad),
'group': ('groups', 1)},
'group': ('groups', 1),
'_target_layout': 'data_layout'},
ignores=['auto_pad'],
custom_check=dimension_constraint())(inputs[:2], attr, params)
use_bias = len(inputs) == 3
if use_bias:
out = _op.nn.bias_add(out, inputs[2])
channel_axis = 3 if attr['_target_layout'] == 'NHWC' else 1
out = _op.nn.bias_add(out, inputs[2], channel_axis)
return out


Expand All @@ -221,6 +241,9 @@ class ConvTranspose(OnnxOpConverter):
"""
@classmethod
def _impl_v1(cls, inputs, attr, params):
if attr['_target_layout'] == 'NHWC':
args['kernel_layout'] = 'HWIO'
inputs[1] = _op.transpose(inputs[1], axes=(2, 3, 1, 0))
# get number of channels
channels = infer_channels(inputs[1], True)
attr['channels'] = channels
Expand All @@ -231,13 +254,15 @@ def _impl_v1(cls, inputs, attr, params):
transforms={
'kernel_shape': 'kernel_size',
'dilations': ('dilation', (0, 0)),
'pads': ('padding', (0, 0), revert_caffe2_pad)
'pads': ('padding', (0, 0), revert_caffe2_pad),
'_target_layout': 'data_layout'
},
disables=['output_shape'],
custom_check=dimension_constraint())(inputs[:2], attr, params)
use_bias = len(inputs) == 3
if use_bias:
out = _op.nn.bias_add(out, inputs[2])
channel_axis = 3 if attr['_target_layout'] == 'NHWC' else 1
out = _op.nn.bias_add(out, inputs[2], channel_axis)
return out


Expand Down Expand Up @@ -322,6 +347,7 @@ class MaxPool(Pool):

@classmethod
def _impl_v8(cls, inputs, attr, params):
attr['storage_order'] = 1 if attr['_target_layout'] == 'NHWC' else attr.get('storage_order', 0)
return AttrCvt(
op_name=dimension_picker(cls.name),
transforms={
Expand Down Expand Up @@ -366,6 +392,9 @@ def _impl_v1(cls, inputs, attr, params):
dims = int(len(pads) / 2)
for i in range(dims):
pad_width.append((pads[i], pads[i+dims]))
if attr['_target_layout'] == 'NHWC' and dims == 4:
pad_width = [pad_width[0], pad_width[2], pad_width[3], pad_width[1]]

attr['pad_width'] = pad_width
pad_mode = attr.get('mode', 'constant').decode('utf-8')
if pad_mode in ['constant', 'edge', 'reflect']:
Expand All @@ -389,6 +418,9 @@ def _impl_v2(cls, inputs, attr, params):
dims = int(len(pads) / 2)
for i in range(dims):
pad_width.append((pads[i], pads[i+dims]))
if attr['_target_layout'] == 'NHWC' and dims == 4:
pad_width = [pad_width[0], pad_width[2], pad_width[3], pad_width[1]]

attr['pad_width'] = pad_width
pad_mode = attr.get('mode', 'constant').decode('utf-8')
if pad_mode in ['constant', 'edge', 'reflect']:
Expand Down Expand Up @@ -443,6 +475,8 @@ class Flatten(OnnxOpConverter):
@classmethod
def _impl_v1(cls, inputs, attr, params):
axis = attr.get('axis', 1)
if attr['_target_layout'] == 'NHWC':
inputs[0] = _op.transpose(inputs[0], axes=(0,3,1,2))
if axis == 1:
out = _op.nn.batch_flatten(inputs[0])
else:
Expand Down Expand Up @@ -548,6 +582,8 @@ class Concat(OnnxOpConverter):

@classmethod
def _impl_v1(cls, inputs, args, params):
axis = args.get('axis', 1)
attr['axis'] = transform_layout_axis(args['_target_layout'], axis)
return AttrCvt(op_name='concatenate')((inputs,), args)

class Scale(OnnxOpConverter):
Expand Down Expand Up @@ -715,6 +751,7 @@ class Unsqueeze(OnnxOpConverter):
@classmethod
def _impl_v1(cls, inputs, attr, params):
for axes in attr['axes']:
axes = transform_layout_axis(attr['_target_layout'], axes)
inputs[0] = _op.expand_dims(inputs[0], axis=axes, num_newaxis=1)
return inputs[0]

Expand Down Expand Up @@ -810,6 +847,7 @@ class Gather(OnnxOpConverter):
@classmethod
def _impl_v1(cls, inputs, attr, params):
axis = attr.get('axis', 0)
axis = transform_layout_axis(attr['_target_layout'], axis)
return AttrCvt('take',
extras={'axis':axis})(inputs, {})

Expand Down Expand Up @@ -900,6 +938,7 @@ class Reduce(OnnxOpConverter):
def _impl_v1(cls, inputs, attr, params):
if 'axes' in attr:
axis = attr.get('axes', 0)
axis = (1,2) if attr['_target_layout'] == 'NHWC' and axis == (2,3) else axis
else:
axis_len = len(infer_shape(inputs[0]))
axis = list(range(axis_len))
Expand Down Expand Up @@ -1222,14 +1261,15 @@ class GraphProto(object):
The input types to the graph
"""

def __init__(self, shape, dtype):
def __init__(self, shape, dtype, layout):
self._nodes = {}
self._params = {}
self._renames = {}
self._num_input = 0
self._num_param = 0
self._shape = shape if shape else {}
self._dtype = dtype
self._layout = layout

def from_onnx(self, graph, opset):
"""Construct Relay expression from ONNX graph.
Expand Down Expand Up @@ -1321,6 +1361,7 @@ def from_onnx(self, graph, opset):
i_name = self._parse_value_proto(node)
attr['tvm_custom'] = {}
attr['tvm_custom']['name'] = i_name
attr['_target_layout'] = self._layout

op = self._convert_operator(op_name, inputs, attr, opset)
node_output = self._fix_outputs(op_name, node.output)
Expand Down Expand Up @@ -1442,6 +1483,7 @@ def _fix_outputs(self, op_name, outputs):
def from_onnx(model,
shape=None,
dtype="float32",
layout="NCHW",
opset=None):
"""Convert a ONNX model into an equivalent Relay Function.

Expand Down Expand Up @@ -1487,7 +1529,7 @@ def from_onnx(model,
warnings.warn(str(e))
except ImportError:
pass
g = GraphProto(shape, dtype)
g = GraphProto(shape, dtype, layout)
graph = model.graph
if opset is None:
try:
Expand Down