def test_kitty_pretrained(self, input_param, image_path, expected_label):
device = "cuda" if torch.cuda.is_available() else "cpu"
# open image
image_size = get_efficientnet_image_size(input_param["model_name"])
img = PIL.Image.open(image_path)
# define ImageNet transforms
tfms = torchvision.transforms.Compose(
[
torchvision.transforms.Resize(image_size),
torchvision.transforms.CenterCrop(image_size),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
]
)
# preprocess and prepare image tensor
img = tfms(img).unsqueeze(0).to(device)
# initialize a pretrained model
net = test_pretrained_networks(EfficientNetBN, input_param, device)
# run inference
with eval_mode(net):
result = net(img)
pred_label = torch.argmax(result, dim=-1)
# check output label
self.assertEqual(pred_label, expected_label)
def test_121_3d_shape_pretrain(self, model, input_param, input_shape,
expected_shape):
net = test_pretrained_networks(model, input_param, device)
net.eval()
with torch.no_grad():
result = net.forward(torch.randn(input_shape).to(device))
self.assertEqual(result.shape, expected_shape)
def test_senet_shape(self, model, input_param):
net = test_pretrained_networks(model, input_param, device)
input_data = torch.randn(3, 3, 64, 64).to(device)
expected_shape = (3, 2)
net = net.to(device)
with eval_mode(net):
result = net(input_data)
self.assertEqual(result.shape, expected_shape)
def test_pretrain_consistency(self, model, input_param, input_shape):
example = torch.randn(input_shape).to(device)
net = test_pretrained_networks(model, input_param, device)
with eval_mode(net):
result = net.features.forward(example)
torchvision_net = torchvision.models.densenet121(
pretrained=True).to(device)
with eval_mode(torchvision_net):
expected_result = torchvision_net.features.forward(example)
self.assertTrue(torch.all(result == expected_result))
def test_pretrain_consistency(self, model, input_param):
input_data = torch.randn(1, 3, 64, 64).to(device)
net = test_pretrained_networks(model, input_param, device)
with eval_mode(net):
result = net.features(input_data)
cadene_net = pretrainedmodels.se_resnet50().to(device)
with eval_mode(cadene_net):
expected_result = cadene_net.features(input_data)
# The difference between Cadene's senet and our version is that
# we use nn.Linear as the FC layer, but Cadene's version uses
# a conv layer with kernel size equals to 1. It may bring a little difference.
self.assertTrue(torch.allclose(result, expected_result, rtol=1e-5, atol=1e-5))
def test_mcfcn_shape(self, input_param, input_shape, expected_shape):
net = test_pretrained_networks(MCFCN, input_param, device)
with eval_mode(net):
result = net.forward(torch.randn(input_shape).to(device))
self.assertEqual(result.shape, expected_shape)
