This tutorial explains how to use pre trained models with PyTorch
.
We will use AlexNet
pre trained model for prediction labels for input image.
Google Colab Notebooks
/content
directory in colab notebook
/content
directory in colab notebook
Install torch
and torchvision
pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html
Instanitate AlexNet
model
import torch
from torchvision import models
from torchvision import transforms
from PIL import Image
alexnet = models.alexnet(pretrained=True)
Pre-process input image for AlexNet
model
preprocess_image = 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]
)])
image = Image.open("./cat3.png")
image_tensor = preprocess_image(image)
Create input tensor from image tensor, by adding one additional dimension
print(image_tensor.shape)
input_tensor = torch.unsqueeze(image_tensor, 0)
print(input_tensor.shape)
Output
torch.Size([3, 224, 224])
torch.Size([1, 3, 224, 224])
Evaluate model and get inference tensor
alexnet.eval()
prediction_tensor = alexnet(input_tensor)
print(prediction_tensor.shape)
Output
torch.Size([1, 1000])
Create list of labels from imagenet_classes
file.
with open('./imagenet_classes.txt') as f:
labels = [line.strip() for line in f.readlines()]
Get index and image label
max_value, index_of_max_value = torch.max(prediction_tensor, 1)
print(index_of_max_value.numpy())
predicted_label = labels[index_of_max_value]
print(predicted_label)
Output
tabby, tabby cat
Category: PyTorch