56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
import torch
|
|
|
|
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
|
|
sys.path.append(MODULE_ROOT)
|
|
from pytorch_wrapper.ptn_max_pool2d import PtnMaxPool2d
|
|
|
|
|
|
class TestPtnMaxPool2dModel(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test instance."""
|
|
self.model_node = PtnMaxPool2d()
|
|
|
|
def test_kernel_stride_padding_as_int(self):
|
|
self.model = self.model_node.f(
|
|
"2", # kernel_size (string)
|
|
"2", # stride (string)
|
|
"0", # padding (string)
|
|
"1" # dilation (string)
|
|
)[0]
|
|
x = torch.ones(1, 3, 64, 64)
|
|
out = self.model(x)
|
|
self.assertEqual(out.size(), torch.Size([1, 3, 32, 32]))
|
|
|
|
def test_kernel_stride_padding_as_tuple(self):
|
|
self.model = self.model_node.f(
|
|
"(2, 2)", # kernel_size (tuple as string)
|
|
"(2, 2)", # stride (tuple as string)
|
|
"(1, 1)", # padding (tuple as string)
|
|
"(1, 1)" # dilation (tuple as string)
|
|
)[0]
|
|
x = torch.ones(1, 3, 64, 64)
|
|
out = self.model(x)
|
|
self.assertEqual(out.size(), torch.Size([1, 3, 33, 33]))
|
|
|
|
def test_padding_as_different_values(self):
|
|
for padding in ["0", "1"]:
|
|
with self.subTest(padding=padding):
|
|
self.model = self.model_node.f(
|
|
"2", # kernel_size (string)
|
|
"2", # stride (string)
|
|
padding, # padding (string)
|
|
"1" # dilation (string)
|
|
)[0]
|
|
x = torch.ones(1, 3, 64, 64)
|
|
out = self.model(x)
|
|
expected_size = 32 + int(padding) # Adjust for padding effect
|
|
self.assertEqual(out.size(), torch.Size([1, 3, expected_size, expected_size]))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|