Backend CUDA example#

This tutorial provides a step-by-step guide on performing inference with the LeNet model using Aidge Backend CUDA. Using randomly generated data as input, we will demonstrate how to pass data through the model and observe its predictions. For the sake of simplicity, we will not train the model. Feel free to replace the onnx model with one already trained.

  • The required libraries for this tutorial are Numpy and PyTorch (optional if you you want to load your own model):

[ ]:

!pip install numpy
!pip install torch

  • We start by creating a LeNet model:

[ ]:

import torch
import torch.nn as nn
import torch.onnx

class LeNet(nn.Module):
    def __init__(self, num_classes):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, kernel_size=2, stride=2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, kernel_size=2, stride=2)
        x = x.view(-1, 16 * 4 * 4)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Instantiate the model
num_classes = 10  # Assuming you're working with MNIST dataset
model = LeNet(num_classes)

# Set the model to evaluation mode
model.eval()

# Example input shape (batch_size, channels, height, width)
dummy_input = torch.randn(1, 1, 28, 28)

# Export the model to ONNX
torch.onnx.export(model, dummy_input, "lenet.onnx", verbose=True)

  • import the needed libraries

[2]:

import aidge_core
import aidge_backend_cuda
import aidge_onnx
import numpy as np

  • load onnx model on Aidge

[ ]:

model = aidge_onnx.load_onnx("lenet.onnx")
aidge_core.remove_flatten(model)

# Configure the model
model.set_datatype(aidge_core.dtype.float32)
model.set_backend("cuda")

  • create input tensor:

[ ]:

# Create an input tensor
input =  np.random.randn(1, 1, 28, 28).astype(np.float32)
input_tensor = aidge_core.Tensor(input)
input_tensor.set_datatype(aidge_core.dtype.float32)
input_tensor.set_backend("cuda")

  • compile the model:

[ ]:

model.compile("cuda", aidge_core.dtype.float32, dims=[input_tensor.dims()])

  • create a scheduler and run inference

[ ]:

# Define the scheduler
scheduler = aidge_core.SequentialScheduler(model)
# Run inference !
scheduler.forward(data=[input_tensor])

  • get the ouput: Before getting the output we need to set it to backend cpu

[ ]:

for outNode in model.get_output_nodes():
    outNode.get_operator().get_output(0).set_backend('cpu')
    output_aidge = np.array(outNode.get_operator().get_output(0))
    print("Aidge output: {}".format(output_aidge))
    # Make sure to set the  output back to "cuda" otherwise the model will not be usable
    outNode.get_operator().get_output(0).set_backend('cuda')

This Page