logo

How to use resnet model in PyTorch

gcptutorials.com PyTorch
Pytorch Hub is a pre-trained model repository. In this post we will see how to use resnet model using Pytorch hub.

Import required libraries and download the pre-trained resnet18 model


import torch
import urllib
from PIL import Image
from torchvision import transforms

model = torch.hub.load('pytorch/vision:v0.9.0', 'resnet152', pretrained=True)

Set the model in evaluation mode.


model.eval()

Download sample image


url, filename = ("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
try:
  urllib.URLopener().retrieve(url, filename)
except:
  urllib.request.urlretrieve(url, filename)

Preprocess the downloaded image to pass as a input to the pre-trained resnet18 model


input_image = Image.open(filename)
preprocess = transforms.Compose([
    transforms.Resize(256),
    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)

Check GPU availability


if torch.cuda.is_available():
    input_batch = input_batch.to('cuda')
    model.to('cuda')

Get the inference with disabling gradient calculation


with torch.no_grad():
    output = model(input_batch)

Get the probabilities using softmax from unnormalized scores in output 


probabilities = torch.nn.functional.softmax(output[0], dim=0)

Download ImageNet labels and store in python list


!wget https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt

with open("imagenet_classes.txt", "r") as f:
  categories = [s.strip() for s in f.readlines()]
print(categories)

Get top 5 probabilities and labels using torch.topk 


top5_prob, top5_catid = torch.topk(probabilities, 5)

for i in range(top5_prob.size(0)):
    print(categories[top5_catid[i]], top5_prob[i].item())

Category: PyTorch