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)