Skip to content
Snippets Groups Projects
Commit 53cd3fd1 authored by auphelia's avatar auphelia
Browse files

[Test] Add unit test for quantized avg pool

parent 987a03da
No related branches found
No related tags found
No related merge requests found
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)
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment