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_func_get_efficientnet_input_shape(self):
for model in get_model_names():
result_shape = get_efficientnet_image_size(model_name=model)
expected_shape = get_expected_model_shape(model)
self.assertEqual(result_shape, expected_shape)
