From 2ac1787e4580521848460215e6b06f4bb1648f06 Mon Sep 17 00:00:00 2001 From: Jordan Gong Date: Sun, 3 Jan 2021 15:08:19 +0800 Subject: Unit testing on auto-encoder, HPM and Part Net --- test/hpm.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 test/hpm.py (limited to 'test/hpm.py') diff --git a/test/hpm.py b/test/hpm.py new file mode 100644 index 0000000..a68337d --- /dev/null +++ b/test/hpm.py @@ -0,0 +1,23 @@ +import torch + +from models import HorizontalPyramidMatching + +T, N, C, H, W = 15, 4, 256, 32, 16 + + +def test_default_hpm(): + hpm = HorizontalPyramidMatching(in_channels=C) + x = torch.rand(T, N, C, H, W) + x = hpm(x) + assert tuple(x.size()) == (1 + 2 + 4, T, N, 128) + + +def test_custom_hpm(): + hpm = HorizontalPyramidMatching(in_channels=2048, + out_channels=256, + scales=(1, 2, 4, 8), + use_avg_pool=True, + use_max_pool=False) + x = torch.rand(T, N, 2048, H, W) + x = hpm(x) + assert tuple(x.size()) == (1 + 2 + 4 + 8, T, N, 256) -- cgit v1.2.3