From 53cd3fd141dcea9cbb31ab53fa154d2eae4be37a Mon Sep 17 00:00:00 2001
From: auphelia <jakobapk@web.de>
Date: Thu, 21 May 2020 13:30:41 +0100
Subject: [PATCH] [Test] Add unit test for quantized avg pool

---
 .../brevitas/test_brevitas_avg_pool_export.py | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)
 create mode 100644 tests/brevitas/test_brevitas_avg_pool_export.py

diff --git a/tests/brevitas/test_brevitas_avg_pool_export.py b/tests/brevitas/test_brevitas_avg_pool_export.py
new file mode 100644
index 000000000..01da19b93
--- /dev/null
+++ b/tests/brevitas/test_brevitas_avg_pool_export.py
@@ -0,0 +1,23 @@
+import onnx  # noqa
+import brevitas.onnx as bo
+from brevitas.nn import QuantAvgPool2d
+import pytest
+
+export_onnx_path = "test_avg_pool.onnx"
+
+
+@pytest.mark.parametrize("kernel_size", [7])
+@pytest.mark.parametrize("stride", [1])
+@pytest.mark.parametrize("signed", [False])
+@pytest.mark.parametrize("bit_width", [4])
+def test_brevitas_avg_pool_export(kernel_size, stride, signed, bit_width):
+    ishape = (1, 1024, 7, 7)
+
+    b_avgpool = QuantAvgPool2d(
+        kernel_size=kernel_size,
+        stride=stride,
+        signed=signed,
+        min_overall_bit_width=bit_width,
+        max_overall_bit_width=bit_width,
+    )
+    bo.export_finn_onnx(b_avgpool, ishape, export_onnx_path)
-- 
GitLab