Skip to content

PyTorch ResNet50

Overview

In this tutorial, we demonstrate how to compile and run inference with a pre-trained ResNet50 model from TorchVision using both the RBLN Python API and PyTorch's torch.compile() API.

Setup & Installation

Before you begin, ensure that your system environment is properly configured and that all required packages are installed. This includes:

  • System Requirements:
  • Packages Requirements:
  • Installation Command:
    1
2
    pip3 install torch torchvision  
pip install --extra-index-url https://pypi.rbln.ai/simple/ rebel-compiler>=0.7.4

Note

RBLN SDK is distributed as a .whl package. Please note that rebel-compiler and rbln_driver require an RBLN Portal account.

Using RBLN Python API

Model Compilation

Load the pre-trained ResNet50 model from TorchVision, set it to evaluation mode, and compile it using the RBLN compiler API.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
import torch  
from torchvision.models import resnet50, ResNet50_Weights  
import rebel  

weights = ResNet50_Weights.DEFAULT  
model = resnet50(weights=weights)  
model.eval()  

compiled_model = rebel.compile_from_torch(  
    model,  
    [('input', [1, 3, 224, 224], torch.float32)],  
)  
compiled_model.save('resnet50.rbln')

Model Inference

Download an example image, preprocess it, load the compiled model using RBLN Runtime, and run inference to display the top predicted class.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
import urllib.request  
from torchvision.io.image import read_image  
from rebel import Runtime  

weights = ResNet50_Weights.DEFAULT

urllib.request.urlretrieve('https://rbln-public.s3.ap-northeast-2.amazonaws.com/images/tabby.jpg', 'input.jpg')  
img = read_image('input.jpg')  
preprocess = weights.transforms()  
batch = preprocess(img).unsqueeze(0)  

module = Runtime('resnet50.rbln', tensor_type='pt')  
output = module(batch)  
_, idx = torch.topk(output, 1, dim=1)  
pred_class = weights.meta['categories'][idx]  
print('Top-1 Predicted Class:', pred_class)

The results will look like this: 

1
Top1 category:  tabby

Using torch.compile() API

Model and Input Preparation

Prepare the ResNet50 model and input image.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
import torch  
from torchvision.models import resnet50, ResNet50_Weights  
import urllib.request  
from torchvision.io.image import read_image  

if torch.__version__ >= '2.5.0':  
    torch._dynamo.config.inline_inbuilt_nn_modules = False  

weights = ResNet50_Weights.DEFAULT  
model = resnet50(weights=weights)  
model.eval()  

img_url = 'https://rbln-public.s3.ap-northeast-2.amazonaws.com/images/tabby.jpg'  
img_path = './tabby.jpg'  
with urllib.request.urlopen(img_url) as response, open(img_path, 'wb') as f:  
    f.write(response.read())  

img = read_image(img_path)  
preprocess = weights.transforms()  
batch = preprocess(img).unsqueeze(0)

Compilation and Inference

Compile the model using torch.compile() with the RBLN backend. The first forward pass triggers the compilation, after which inference is run on the compiled model.

Note

To use the RBLN backend with torch.compile(), you must specify backend='rbln'. This requires import rebel before compilation.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
import rebel  # Need to use torch dynamo's "rbln" backend.

compiled_model = torch.compile(  
    model,  
    backend='rbln',  
    options={'cache_dir': './.cache'},  
    dynamic=False  
)  

# Trigger compilation with the first forward pass  
compiled_model(batch)  

rbln_result = compiled_model(batch)  
score, class_id = torch.topk(rbln_result, 1, dim=1)  
category_name = weights.meta['categories'][class_id]  
print('Top-1 category:', category_name)

The results will look like this: 

1
Top1 category:  tabby

Summary and References

This tutorial demonstrated how to compile and run inference on the TorchVision ResNet50 model using both the RBLN Python API and PyTorch's torch.compile() API. You can now deploy the compiled model on an RBLN NPU for efficient inference.

References: