Use your AI Model with Raspberry PI and PyTorch for Image Classification

Last Updated on 13th September 2025 by peppe8o

This tutorial will show you how to use Artificial Intelligence on Raspberry PI with PyTorch to perform image classification tasks. This is part of a series of AI tutorials for Raspberry PI, and you can find the starting point in my Create your Artificial Intelligence (AI) model with Raspberry PI and PyTorch for Image Classification article, where I explained how to create your Artificial Intelligence model.

For this reason, this article will start from the model saved as my_model.pth from the previous guide.

The Neural Network Model

It is important to note that the model, as defined, requires input images of 32×32 pixels. The code described in this tutorial will allow you to resize images before processing them, as you can use the model with images of whatever size.

When you use an AI model for predictions, you must use the same neural network structure defined in the training process, as the saved model includes the node parameters for that network configuration.

So, respecting the model defined in the previous article, our neural network will be composed by:

• A convolution layer (with 3 input channels and 8 output channels)
• A maxpool layer (with 8 input channels and 8 output channels)
• A flatting operation converting the 8 input channels into 2048 nodes, and a linear layer (with 2048 input nodes and 10 output channels)

The previous article shows how these layers work in the forward process, and the model evaluation process will also use this during the predictions. I suggest you visit it to get a base knowledge of the network structure.

What We Need

As usual, I suggest adding from now to your favourite e-commerce shopping cart all the needed hardware, so that at the end you will be able to evaluate overall costs and decide if to continue with the project or remove them from the shopping cart. So, hardware will be only:

• Raspberry PI Computer Board (including proper power supply or using a smartphone micro USB charger with at least 3A)
• high-speed micro SD card (at least 16 GB, at least class 10)
• (optional, but suggested) an active cooling for Raspberry PI. You can see the benefits in my Active Cooling for Raspberry PI article

For this tutorial, I will use my Raspberry PI 5 Model B with 8GB of RAM. If you use a Raspberry PI model with a lower amount of RAM, I strongly suggest increasing the Raspberry PI SWAP memory with the linked tutorial.

Step-by-step Procedure

In the previous article, I showed you how to install the required packages into a virtual environment. To run the prediction script we need to activate the same Python environment:

source ./my_ai/bin/activate

With this, you can now download my test script directly from my download area:

wget https://peppe8o.com/download/python/pytorch/predict.py

The following paragraphs will explain the code line by line, before showing you the results.

Import the Required Packages and the Neural Network Structure

At the beginning, we need to import the required packages. The sys package isn’t required for the model execution: I use it to enable you to give the input image from the bash shell (as an argument), without the need to edit the script at every image you want to analyse. We’ll see it in a few lines.

import torch, torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import sys

Then, we must define the neural network. Here, we must use the same structure as the network used for training, and this has been explained in my previous article:

class NeuralNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 8, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(8 * 16 * 16, 10)

    def forward(self, x):
        # Layer 1
        x = self.conv1(x)
        x = F.relu(x)
        # Layer 2
        x = self.pool(x)
        # Layer 3
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)

        return x

We can now define the classes for the labels. Also, these must be the same (and in the same order) as those defined in the train process:

classes = ["airplane", "car", "bird", "cat", "deer",
           "dog", "frog", "horse", "ship", "truck"]

Our initial part of the script defines the transform function, which will manage the image that is obtained from the user for evaluation. The transform process will resize the image to 32×32 pixels, so that it will fit the sizes required by the model, converted to tensors and then normalised as explained in the train model:

transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5),
                         (0.5, 0.5, 0.5))
])

Initialise the AI Model

At this point, we can create the model according to the given class:

net = NeuralNet()

The parameters for this model have been saved in our my_model.pth file, so we can import it:

PATH = './my_model.pth'
net.load_state_dict(torch.load(PATH))

At the end of this part, we must set our AI model into evaluation mode:

net.eval()

Prepare the Image to Classify

The images coming for classification can have several different sizes. We must prepare them to be analysed by our AI model. Our neural network has been trained with input models of 32×32 pixels and in RGB. To make the image compliant with this, we’ll use the PIL (Python Imaging Library) library. Before opening the image, we’ll import this library:

from PIL import Image

Instead of providing the image link inside the script, which would require you to edit the file to update the image path for every image you want to analyse, in this script we’ll use the sys.argv method, which will enable you to provide the image path when you run the Python script from the terminal shell. We’ll see how to use this in the final chapter of this post. Alternatively, you can set a fixed image path to the image_path variable.

In the following lines, we get the image path, open the image, and convert it to RGB:

image_path = sys.argv[1]
img = Image.open(image_path).convert("RGB")

Finally, we use the transform to adjust the image size and convert it to tensors:

img_tensor = transform(img).unsqueeze(0)

Analyse the Image With your AI Model and Get Predictions

Finally, we can perform the inference with the image:

outputs = net(img_tensor)

As I want to get the confidence evaluation for each prediction, we’ll use the Softmax function to normalise the result tensors so that we can convert the resulting tensors to percentage values. The new “probs” variable will include both the confidence percentages as well as the labels:

probs = torch.nn.functional.softmax(outputs, dim=1)

We’ll explore 2 different methods to get the classification predictions from our AI model. In the first method, we’ll get the label with the maximum probability to match the image classification. With the second method, we’ll get several predictions, ordered by probability.

For the first method, we’ll initially print a small banner:

print("---------------------")
print("Max probability")
print("---------------------")

The torch.max The function will extract from the results the one with the maximum probability to match the image classification. It will return both the confidence, in percentage, and the label ID, which will allow us to retrieve the classification label from the classes object:

conf, label = torch.max(probs, 1)

Finally, we print the results:

print(f"Prediction: {classes[label.item()]}, Confidence: {conf[0]:.2%}")

For the second method, again, we print a small banner:

print("---------------------")
print("Top probabilities")
print("---------------------")

With the torch.topk function we can retrieve a defined number of matching cases ordered by confidence. We can adjust the number of results by setting it in the function options. For this test, I’ve set it to 4:

conf, label = torch.topk(probs,4)

The final loop will print all the results, each one with its confidence percentage:

for p in range(label.numel()):
  print(f"Prediction: {classes[label[0][p].item()]}, Confidence: {conf[0][p].item():.2%}")

Use the Artificial Intelligence Model on Raspberry PI

You can run the model with the terminal command:

python predict.py image.jpg

As you can see, you can give the path for the image to analyse. In my case, the “image.jpg” was an image from an airplane, saved in the folder where the command runs.

This will give you a result similar to the following:

(my_ai) pi@raspberrypi:~ $ python predict.py image.jpg
---------------------
Max probability
---------------------
Prediction: airplane, Confidence: 66.48%

---------------------
Top probabilities
---------------------
Prediction: airplane, Confidence: 66.48%
Prediction: ship, Confidence: 15.99%
Prediction: bird, Confidence: 5.91%
Prediction: car, Confidence: 4.16%

What’s Next

If you are interested in more Raspberry PI projects (both with Lite and Desktop OS), take a look at my Raspberry PI Tutorials.

Enjoy!