def test_ill_input_shape(self):
net = UNet(spatial_dims=2,
in_channels=1,
out_channels=3,
channels=(16, 32, 64),
strides=(2, 2))
with eval_mode(net):
with self.assertRaisesRegex(RuntimeError,
"Sizes of tensors must match"):
net.forward(torch.randn(2, 1, 16, 5))
def test_shape(self, input_param, input_shape, expected_shape):
net = UNet(**input_param).to(device)
net.eval()
with torch.no_grad():
result = net.forward(torch.randn(input_shape).to(device))
self.assertEqual(result.shape, expected_shape)
def test_shape(self, input_param, input_data, expected_shape):
net = UNet(**input_param)
net.eval()
with torch.no_grad():
result = net.forward(input_data)
self.assertEqual(result.shape, expected_shape)