65 lines
2.0 KiB
Python
65 lines
2.0 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_conv_model import PtnConvModel
|
|
|
|
|
|
class TestPtnConvModel(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Set up test instance."""
|
|
self.model_node = PtnConvModel()
|
|
|
|
def test_forward(self):
|
|
model = self.model_node.f(
|
|
input_dim="(4, 28, 28)",
|
|
penultimate_dim=0,
|
|
output_dim=10,
|
|
channel_list="[32,64,128,256]",
|
|
kernel_size_list="[3,3,3,3]",
|
|
padding_list="[1,1,1,1]",
|
|
downsample_list="[True,True,True,True]"
|
|
)[0]
|
|
bs = 2
|
|
x = torch.ones(bs, 4, 28, 28)
|
|
out = model(x)
|
|
self.assertEqual(out.size(), torch.Size([bs, 10]))
|
|
|
|
def test_rank3_input(self):
|
|
model = self.model_node.f(
|
|
input_dim="(1, 28, 28)",
|
|
penultimate_dim=0,
|
|
output_dim=10,
|
|
channel_list="[32,64,128,256]",
|
|
kernel_size_list="[3,3,3,3]",
|
|
padding_list="[1,1,1,1]",
|
|
downsample_list="[True,True,True,True]"
|
|
)[0]
|
|
bs = 2
|
|
x = torch.ones(bs, 28, 28)
|
|
out = model(x)
|
|
self.assertEqual(out.size(), torch.Size([bs, 10]))
|
|
|
|
def test_rank2_input(self):
|
|
model = self.model_node.f(
|
|
input_dim="(1, 28, 28)",
|
|
penultimate_dim=0,
|
|
output_dim=10,
|
|
channel_list="[32,64,128,256]",
|
|
kernel_size_list="[3,3,3,3]",
|
|
padding_list="[1,1,1,1]",
|
|
downsample_list="[True,True,True,True]"
|
|
)[0]
|
|
bs = 2
|
|
x = torch.ones(bs, 28)
|
|
with self.assertRaises(ValueError) as context:
|
|
model(x)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|