159 lines
5.3 KiB
Python
159 lines
5.3 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_def import ConvModel, ConvSingleLayer
|
|
|
|
|
|
class TestPtnConvModelDef(unittest.TestCase):
|
|
def setUp(self):
|
|
"""Setup method for initializing test conditions."""
|
|
pass
|
|
|
|
def test_single_layer_1_with_padding(self):
|
|
"""Test ConvSingleLayer with padding=1, no downsampling."""
|
|
model = ConvSingleLayer(
|
|
in_channels=4,
|
|
out_channels=32,
|
|
kernel_size=3,
|
|
padding=1,
|
|
downsample=False,
|
|
apply_activation=True)
|
|
x = torch.ones(1, 4, 28, 28)
|
|
out = model(x)
|
|
self.assertEqual(out.size(), torch.Size([1, 32, 28, 28]))
|
|
|
|
def test_single_layer_no_padding(self):
|
|
"""Test ConvSingleLayer with padding=0, no downsampling."""
|
|
model = ConvSingleLayer(
|
|
in_channels=4,
|
|
out_channels=32,
|
|
kernel_size=3,
|
|
padding=0,
|
|
downsample=False,
|
|
apply_activation=True)
|
|
x = torch.ones(1, 4, 28, 28)
|
|
out = model(x)
|
|
self.assertEqual(out.size(), torch.Size([1, 32, 26, 26]))
|
|
|
|
def test_single_layer_3_no_padding_downsample(self):
|
|
"""Test ConvSingleLayer with padding=0 and downsampling enabled."""
|
|
model = ConvSingleLayer(
|
|
in_channels=4,
|
|
out_channels=32,
|
|
kernel_size=3,
|
|
padding=0,
|
|
downsample=True,
|
|
apply_activation=True)
|
|
x = torch.ones(1, 4, 28, 28)
|
|
out = model(x)
|
|
self.assertEqual(out.size(), torch.Size([1, 32, 13, 13]))
|
|
|
|
def test_1(self):
|
|
"""Test ConvModel with a single convolutional layer, checking output size."""
|
|
model = ConvModel(
|
|
input_dim=(4, 28, 28),
|
|
penultimate_dim=0,
|
|
output_dim=10,
|
|
channel_list=[32],
|
|
kernel_size_list=[3],
|
|
padding_list=[1],
|
|
downsample_list=[False])
|
|
|
|
x = torch.ones(1, 4, 28, 28)
|
|
out = model(x)
|
|
self.assertEqual(out.size(), torch.Size([1, 10]))
|
|
|
|
def test_2(self):
|
|
"""Test ConvModel intermediate convolution layers' output size."""
|
|
model = ConvModel(
|
|
input_dim=(4, 28, 28),
|
|
penultimate_dim=0,
|
|
output_dim=10,
|
|
channel_list=[32],
|
|
kernel_size_list=[3],
|
|
padding_list=[1],
|
|
downsample_list=[False])
|
|
|
|
x = torch.ones(1, 4, 28, 28)
|
|
for c in model.conv_layers:
|
|
x = c(x)
|
|
out = x
|
|
self.assertEqual(out.size(), torch.Size([1, 32, 28, 28]))
|
|
|
|
def test_3(self):
|
|
"""Test ConvModel with two convolutional layers and downsampling enabled."""
|
|
model = ConvModel(
|
|
input_dim=(4, 28, 28),
|
|
penultimate_dim=0,
|
|
output_dim=10,
|
|
channel_list=[32,64],
|
|
kernel_size_list=[3,3],
|
|
padding_list=[1,1],
|
|
downsample_list=[True,True])
|
|
|
|
x = torch.ones(1, 4, 28, 28)
|
|
for c in model.conv_layers:
|
|
x = c(x)
|
|
out = x
|
|
self.assertEqual(out.size(), torch.Size([1, 64, 7, 7]))
|
|
|
|
def test_4(self):
|
|
"""Test ConvModel with three convolutional layers and downsampling enabled."""
|
|
model = ConvModel(
|
|
input_dim=(4, 28, 28),
|
|
penultimate_dim=0,
|
|
output_dim=10,
|
|
channel_list=[32,64,128],
|
|
kernel_size_list=[3,3,3],
|
|
padding_list=[1,1,1],
|
|
downsample_list=[True,True,True])
|
|
|
|
x = torch.ones(1, 4, 28, 28)
|
|
for c in model.conv_layers:
|
|
x = c(x)
|
|
out = x
|
|
self.assertEqual(out.size(), torch.Size([1, 128, 3, 3]))
|
|
|
|
def test_5(self):
|
|
"""Test ConvModel with four convolutional layers and downsampling enabled."""
|
|
model = ConvModel(
|
|
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])
|
|
|
|
x = torch.ones(1, 4, 28, 28)
|
|
for c in model.conv_layers:
|
|
x = c(x)
|
|
out = x
|
|
self.assertEqual(out.size(), torch.Size([1, 256, 1, 1]))
|
|
|
|
def test_6(self):
|
|
"""Test ConvModel with five convolutional layers, last layer without downsampling."""
|
|
model = ConvModel(
|
|
input_dim=(4, 28, 28),
|
|
penultimate_dim=0,
|
|
output_dim=10,
|
|
channel_list=[32,64,128,256,512],
|
|
kernel_size_list=[3,3,3,3,1],
|
|
padding_list=[1,1,1,1,0],
|
|
downsample_list=[True,True,True,True,False])
|
|
|
|
x = torch.ones(1, 4, 28, 28)
|
|
for c in model.conv_layers:
|
|
x = c(x)
|
|
out = x
|
|
self.assertEqual(out.size(), torch.Size([1, 512, 1, 1]))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|