ShuffleNet-v2-Pytorch
ShuffleNet-v2-Pytorch copied to clipboard
Inference example
Hi,
Thanks for your work. I trained a model with my own data. How do I infer a single image? I wrote this script but I get wrong result
import torch
from models import ShuffleNet2
num_classes = 2
input_size = 224
net_type = 1
model = ShuffleNet2(num_classes, input_size, net_type)
model_path = "./save/ShuffleNet21660477670.pkl"
model.load_state_dict(torch.load(model_path))
model.to('cuda')
model.eval()
from PIL import Image
from torchvision import transforms
input_image = Image.open('./dogvscat/train/cat.1.jpg')
preprocess = transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0) # create a mini-batch as expected by the model
# move the input and model to GPU for speed if available
if torch.cuda.is_available():
input_batch = input_batch.to('cuda')
model.to('cuda')
with torch.no_grad():
output = model(input_batch)
print(output[0])
# The output has unnormalized scores. To get probabilities, you can run a softmax on it.
probabilities = torch.nn.functional.softmax(output[0], dim=0)
print(probabilities)
# Show top categories per image
top5_prob, top5_catid = torch.topk(probabilities, 2)
for i in range(top5_prob.size(0)):
print(top5_catid[i], top5_prob[i].item())