def __init__(self,
in_channels,
out_channels,
dropout_p,
spatial_dims: int = 2):
super().__init__()
self.conv_conv_se = nn.Sequential(
Convolution(spatial_dims,
in_channels,
out_channels,
kernel_size=3,
norm=Norm.BATCH,
act=Act.LEAKYRELU),
nn.Dropout(dropout_p),
Convolution(spatial_dims,
out_channels,
out_channels,
kernel_size=3,
norm=Norm.BATCH,
act=Act.LEAKYRELU),
ResidualSELayer(spatial_dims=spatial_dims,
in_channels=out_channels,
r=2),
)
def test_shape(self, input_param, input_data, expected_shape):
net = ResidualSELayer(**input_param)
net.eval()
with torch.no_grad():
result = net(input_data)
self.assertEqual(result.shape, expected_shape)
def test_shape(self, input_param, input_shape, expected_shape):
net = ResidualSELayer(**input_param)
with eval_mode(net):
result = net(torch.randn(input_shape))
self.assertEqual(result.shape, expected_shape)
def test_script(self):
input_param, input_shape, _ = TEST_CASES[0]
net = ResidualSELayer(**input_param)
test_data = torch.randn(input_shape)
test_script_save(net, test_data)
def test_script(self):
input_param, input_shape, _ = TEST_CASES[0]
net = ResidualSELayer(**input_param)
test_data = torch.randn(input_shape)
out_orig, out_reloaded = test_script_save(net, test_data)
assert torch.allclose(out_orig, out_reloaded)