68 lines
1.9 KiB
Python
68 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_resnet_model import PtnResnetModel
|
|
|
|
|
|
class TestPtnResnetModel(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test instance."""
|
|
self.model_node = PtnResnetModel()
|
|
|
|
def test_forward(self):
|
|
model = self.model_node.f(
|
|
input_dim="(4, 28, 28)",
|
|
output_dim=10,
|
|
num_blocks=2
|
|
)[0]
|
|
bs = 2
|
|
x = torch.ones(bs, 4, 28, 28)
|
|
out = model(x)
|
|
self.assertEqual(out.size(), torch.Size([bs, 10]))
|
|
|
|
def test_forward2(self):
|
|
model = self.model_node.f(
|
|
input_dim="(3, 256, 256)",
|
|
output_dim=4,
|
|
num_blocks=2
|
|
)[0]
|
|
bs = 4
|
|
x = torch.ones(bs, 3, 256, 256)
|
|
out = model(x)
|
|
self.assertEqual(out.size(), torch.Size([bs, 4]))
|
|
|
|
def test_invalid_input_shape(self):
|
|
model = self.model_node.f(
|
|
input_dim="(3, 32, 32)",
|
|
output_dim=10,
|
|
num_blocks=2
|
|
)[0]
|
|
bs = 2
|
|
x = torch.ones(bs, 32, 32) # Missing channel dimension
|
|
with self.assertRaises(RuntimeError) as context:
|
|
model(x)
|
|
|
|
def test_invalid_output_dim(self):
|
|
with self.assertRaises(ValueError):
|
|
self.model_node.f(
|
|
input_dim="(3, 28, 28)",
|
|
output_dim=0, # Invalid output_dim
|
|
num_blocks=2
|
|
)
|
|
|
|
def test_invalid_num_blocks(self):
|
|
with self.assertRaises(ValueError):
|
|
self.model_node.f(
|
|
input_dim="(3, 28, 28)",
|
|
output_dim=10,
|
|
num_blocks=0 # Invalid num_blocks
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|