Skip to content

Commit

Permalink
fix prelu importer and add tests: (apache#5521)
Browse files Browse the repository at this point in the history
  • Loading branch information
Matthew Brookhart authored and Trevor Morris committed Jun 18, 2020
1 parent 72628c6 commit d8e65c4
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 1 deletion.
7 changes: 6 additions & 1 deletion python/tvm/relay/frontend/onnx.py
Original file line number Diff line number Diff line change
Expand Up @@ -576,7 +576,12 @@ class Prelu(OnnxOpConverter):
@classmethod
def _impl_v1(cls, inputs, attr, params):
assert len(inputs) == 2, "Prelu need 2 inputs, {} given".format(len(inputs))
return _op.nn.prelu(inputs[0], inputs[1])
alpha_shape = infer_shape(inputs[1])
if len(alpha_shape) != 1:
alpha = _op.reshape(inputs[1], (-1,))
else:
alpha = inputs[1]
return _op.nn.prelu(inputs[0], alpha)


class Reciprocal(OnnxOpConverter):
Expand Down
29 changes: 29 additions & 0 deletions tests/python/frontend/onnx/test_forward.py
Original file line number Diff line number Diff line change
Expand Up @@ -1532,6 +1532,34 @@ def selu_x(x, alpha, gamma):
{'alpha': 0.25, 'gamma': 0.3})


def test_prelu():
def verify_prelu(x_shape, a_shape):
node = helper.make_node('PRelu',
inputs=['X', 'slope'],
outputs=['Y'])

graph = helper.make_graph([node],
"prelu_test",
inputs=[helper.make_tensor_value_info("X", TensorProto.FLOAT, list(x_shape)),
helper.make_tensor_value_info("slope", TensorProto.FLOAT, list(a_shape))],
outputs=[helper.make_tensor_value_info("Y", TensorProto.FLOAT, list(x_shape))])

model = helper.make_model(graph, producer_name='prelu_test')

indata = np.random.uniform(-10, 10, x_shape).astype(np.float32)
slopedata = np.random.uniform(-10, 10, a_shape).astype(np.float32)
onnx_out = get_onnxruntime_output(model, [indata, slopedata])

for target, ctx in [('llvm', tvm.cpu())]:
tvm_out = get_tvm_output(model, [indata, slopedata], target, ctx, list(x_shape),
output_dtype='float32')
tvm.testing.assert_allclose(onnx_out[0], tvm_out, rtol=1e-05, atol=1e-05)

verify_prelu([3,4,5,6], [1, 4, 1, 1])
verify_prelu([1,8,5,6], [1, 8, 1, 1])
verify_prelu([2,12,16,16], [1, 12, 1, 1])


def test_ThresholdedRelu():
def ThresholdedRelu_x(x, alpha):
out_np = np.clip(x, alpha, np.inf)
Expand Down Expand Up @@ -2535,6 +2563,7 @@ def verify_roi_align(input_dims, num_roi, output_height, output_width, sampling_
test_leaky_relu()
test_elu()
test_selu()
test_prelu()
test_ThresholdedRelu()
test_ScaledTanh()
test_ParametricSoftplus()
Expand Down

0 comments on commit d8e65c4

Please sign in to comment.