From 4e494317e3dc50436bee11b55f1b34d0c700b530 Mon Sep 17 00:00:00 2001 From: Mirza Mrahorovic <34712307+mmrahorovic@users.noreply.github.com> Date: Sat, 20 Feb 2021 18:09:11 +0100 Subject: [PATCH] [batchnorm_to_affine]: epsilon value is now read out from the attributes. (#21) [test_batchnorm_to_affine]: added a test case for various epsilon values. --- .../transformation/batchnorm_to_affine.py | 4 +- .../test_batchnorm_to_affine.py | 53 +++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/src/finn/transformation/batchnorm_to_affine.py b/src/finn/transformation/batchnorm_to_affine.py index 3917e78..5204d07 100644 --- a/src/finn/transformation/batchnorm_to_affine.py +++ b/src/finn/transformation/batchnorm_to_affine.py @@ -32,6 +32,7 @@ from finn.transformation.base import Transformation from finn.transformation.infer_shapes import InferShapes +from finn.util.basic import get_by_name class BatchNormToAffine(Transformation): @@ -52,7 +53,8 @@ def apply(self, model): bias = model.get_initializer(n.input[2]) mean = model.get_initializer(n.input[3]) variance = model.get_initializer(n.input[4]) - epsilon = 1e-5 + epsilon = get_by_name(n.attribute, "epsilon") + epsilon = getattr(epsilon, "f", 1e-5) # find A and B to compute batchnorm as affine transpose Ax+B # TODO is a division by moving avg factor needed for variance? A = scale / np.sqrt(epsilon + variance) diff --git a/tests/transformation/test_batchnorm_to_affine.py b/tests/transformation/test_batchnorm_to_affine.py index 984a996..09e4fc8 100644 --- a/tests/transformation/test_batchnorm_to_affine.py +++ b/tests/transformation/test_batchnorm_to_affine.py @@ -29,6 +29,7 @@ import pytest import numpy as np +import onnx import os import urllib.request as ureq @@ -65,3 +66,55 @@ def test_batchnorm_to_affine_shufflenet(): produced = oxe.execute_onnx(new_model, input_dict)[oname] assert np.isclose(expected, produced).all() os.remove(export_onnx_path) + + +@pytest.mark.parametrize("epsilon", [0.0, 0.00001, 0.001]) +def test_batchnorm_to_affine_epsilon(epsilon): + """ Dummy batchnorm node to test out the epsilon attribute. """ + + batchnorm_node = onnx.helper.make_node( + "BatchNormalization", + inputs=["x", "s", "bias", "mean", "var"], + outputs=["y"], + epsilon=epsilon, + ) + + x = onnx.helper.make_tensor_value_info("x", onnx.TensorProto.FLOAT, [1, 3, 5, 5]) + s = onnx.helper.make_tensor_value_info("s", onnx.TensorProto.FLOAT, [3]) + bias = onnx.helper.make_tensor_value_info("bias", onnx.TensorProto.FLOAT, [3]) + mean = onnx.helper.make_tensor_value_info("mean", onnx.TensorProto.FLOAT, [3]) + var = onnx.helper.make_tensor_value_info("var", onnx.TensorProto.FLOAT, [3]) + y = onnx.helper.make_tensor_value_info("y", onnx.TensorProto.FLOAT, [1, 3, 5, 5]) + + # Graph + graph = onnx.helper.make_graph( + nodes=[batchnorm_node], + name="test_batchnorm_graph", + inputs=[x], + outputs=[y], + value_info=[s, bias, mean, var], + ) + + onnx_model = onnx.helper.make_model(graph, producer_name="test_batchnorm-model") + model = ModelWrapper(onnx_model) + + model.set_initializer("s", np.array([1, 2, 3]).astype(np.float32)) + model.set_initializer("bias", np.array([1, 2, 3]).astype(np.float32)) + model.set_initializer("mean", np.array([3, 4, 5]).astype(np.float32)) + model.set_initializer("var", np.array([0.5, 0.7, 0.3]).astype(np.float32)) + + i_val = np.arange(0, 3 * 5 * 5, dtype=np.float32) + i_val = np.reshape(i_val, [1, 3, 5, 5]) + input_dict = {"x": i_val} + output_node_name = "y" + + output_dict = oxe.execute_onnx(model, input_dict, return_full_exec_context=True) + output_original = output_dict[output_node_name] + + model_lowered = model.transform(BatchNormToAffine()) + output_dict = oxe.execute_onnx( + model_lowered, input_dict, return_full_exec_context=True + ) + output_lowered = output_dict[output_node_name] + + assert (output_original == output_lowered).all()