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
```

Categories: PyTorch

Similar Articles