Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Commit

Permalink
Handle default p_value
Browse files Browse the repository at this point in the history
  • Loading branch information
vandanavk committed Dec 21, 2018
1 parent 6d7ddf5 commit 703bac8
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 5 deletions.
9 changes: 6 additions & 3 deletions python/mxnet/contrib/onnx/mx2onnx/_op_translations.py
Original file line number Diff line number Diff line change
Expand Up @@ -586,7 +586,7 @@ def convert_pooling(node, **kwargs):
pool_type = attrs["pool_type"]
stride = eval(attrs["stride"]) if attrs.get("stride") else None
global_pool = get_boolean_attribute_value(attrs, "global_pool")
p_value = int(attrs.get('p_value', '2'))
p_value = attrs.get('p_value', 'None')

pooling_convention = attrs.get('pooling_convention', 'valid')

Expand All @@ -603,13 +603,16 @@ def convert_pooling(node, **kwargs):
global_pool_types = {"max": "GlobalMaxPool", "avg": "GlobalAveragePool",
"lp": "GlobalLpPool"}

if pool_type == 'lp' and p_value == 'None':
raise AttributeError('ONNX requires a p value for LpPool and GlobalLpPool')

if global_pool:
if pool_type == 'lp':
node = onnx.helper.make_node(
global_pool_types[pool_type],
input_nodes, # input
[name],
p=p_value,
p=int(p_value),
name=name
)
else:
Expand All @@ -625,7 +628,7 @@ def convert_pooling(node, **kwargs):
pool_types[pool_type],
input_nodes, # input
[name],
p=p_value,
p=int(p_value),
kernel_shape=kernel,
pads=pad_dims,
strides=stride,
Expand Down
44 changes: 42 additions & 2 deletions tests/python-pytest/onnx/export/mxnet_export_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,8 +115,11 @@ def forward_pass(sym, arg, aux, data_names, input_data):
data_forward.append(mx.nd.array(val))

mod.bind(for_training=False, data_shapes=data_shapes, label_shapes=None)
mod.set_params(arg_params=arg, aux_params=aux,
allow_missing=True, allow_extra=True)
if not arg and not aux:
mod.init_params()
else:
mod.set_params(arg_params=arg, aux_params=aux,
allow_missing=True, allow_extra=True)

# run inference
batch = namedtuple('Batch', ['data'])
Expand Down Expand Up @@ -409,6 +412,43 @@ def test_ops(op_name, inputs, input_tensors, numpy_op):
np.logical_not(input_data[0]).astype(np.float32))


@with_seed()
def testLpPooling():
def test_pooling(opname, data, attrs, p):
input1 = np.random.rand(*data).astype("float32")
inputs = [helper.make_tensor_value_info("input1", TensorProto.FLOAT, shape=data)]
sym = mx.sym.Pooling(mx.sym.Variable('input1'), pool_type='lp', p_value=p, **attrs)
lppool_output = forward_pass(sym, None, None, ['input1'], input1)

lppool_op_tensor = [helper.make_tensor_value_info("output", TensorProto.FLOAT, shape=np.shape(lppool_output))]

if attrs.get('global_pool', False):
lppool_node = [helper.make_node(opname, ["input1"], ["output"], p=p)]
else:
lppool_node = [helper.make_node(opname, ["input1"], ["output"], p=p, **attrs)]

lppool_graph = helper.make_graph(lppool_node,
opname+"_test",
inputs,
lppool_op_tensor)

lppool_model = helper.make_model(lppool_graph)

bkd_rep = backend.prepare(lppool_model)
output = bkd_rep.run([input1])

npt.assert_almost_equal(output[0], lppool_output)

ip = (2, 3, 20, 20)
kernel = (4, 5)
pad = (0, 0)
stride = (1, 1)

for p_value in range(1, 3):
test_pooling('LpPool', ip, {'kernel': kernel, 'stride': stride, 'pad': pad}, p=p_value)
test_pooling('GlobalLpPool', ip, {'kernel': kernel, 'stride': stride, 'pad': pad, 'global_pool': True}, p=p_value)


def _assert_sym_equal(lhs, rhs):
assert lhs.list_inputs() == rhs.list_inputs() # input names must be identical
assert len(lhs.list_outputs()) == len(rhs.list_outputs()) # number of outputs must be identical
Expand Down

0 comments on commit 703bac8

Please sign in to comment.