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)