def test_script(self):
net = EfficientNetBN(model_name="efficientnet-b0",
spatial_dims=2,
in_channels=3,
num_classes=1000)
net.set_swish(
memory_efficient=False
) # at the moment custom memory efficient swish is not exportable with jit
test_data = torch.randn(1, 3, 224, 224)
test_script_save(net, test_data)
def test_shape(self, input_param, input_shape, expected_shape):
device = "cuda" if torch.cuda.is_available() else "cpu"
with skip_if_downloading_fails():
net = EfficientNetBN(**input_param).to(device)
# run inference with random tensor
with eval_mode(net):
result = net(torch.randn(input_shape).to(device))
# check output shape
self.assertEqual(result.shape, expected_shape)
def test_shape(self, input_param, input_shape, expected_shape):
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
# initialize model
net = EfficientNetBN(**input_param).to(device)
except (ContentTooShortError, HTTPError, RuntimeError) as e:
print(str(e))
return # skipping the tests because of http errors
# run inference with random tensor
with eval_mode(net):
result = net(torch.randn(input_shape).to(device))
# check output shape
self.assertEqual(result.shape, expected_shape)
def test_non_default_shapes(self, input_param, input_shape, expected_shape):
device = "cuda" if torch.cuda.is_available() else "cpu"
with skip_if_downloading_fails():
net = EfficientNetBN(**input_param).to(device)
# override input shape with different variations
num_dims = len(input_shape) - 2
non_default_sizes = [128, 256, 512]
for candidate_size in non_default_sizes:
input_shape = input_shape[0:2] + (candidate_size,) * num_dims
# run inference with random tensor
with eval_mode(net):
result = net(torch.randn(input_shape).to(device))
# check output shape
self.assertEqual(result.shape, expected_shape)
def test_non_default_shapes(self, input_param, input_shape, expected_shape):
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
# initialize model
net = EfficientNetBN(**input_param).to(device)
except (ContentTooShortError, HTTPError, RuntimeError) as e:
print(str(e))
return # skipping the tests because of http errors
# override input shape with different variations
num_dims = len(input_shape) - 2
non_default_sizes = [128, 256, 512]
for candidate_size in non_default_sizes:
input_shape = input_shape[0:2] + (candidate_size,) * num_dims
# run inference with random tensor
with eval_mode(net):
result = net(torch.randn(input_shape).to(device))
# check output shape
self.assertEqual(result.shape, expected_shape)
def test_ill_arg(self):
with self.assertRaises(ValueError):
# wrong spatial_dims
EfficientNetBN(model_name="efficientnet-b0", spatial_dims=4)
# wrong model_name
EfficientNetBN(model_name="efficientnet-b10", spatial_dims=3)