Skip to content

PyTorch ResNet50

Overview

This tutorial shows 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:
  • Package Requirements:
  • Installation Command:
    1
2
3
4
    pip install --extra-index-url https://download.pytorch.org/whl/cpu torch torchvision
pip install \
  --extra-index-url https://pypi.rbln.ai/simple \
  rebel-compiler==0.12.0

Info

  • rebel-compiler requires an RBLN Portal account.
  • The commands above are intended for a default pip install on Debian-based Linux such as Ubuntu. For all other configurations, refer to the Installation Guide for the supported install matrix and the applicable commands.

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.

 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
17
18
import urllib.request  
from torchvision.io.image import read_image  
import torch  
from torchvision.models import ResNet50_Weights  
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)

Example Output: 

1
Top-1 Predicted Class: tabby

Using the 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
import torch  
from torchvision.models import resnet50, ResNet50_Weights  
import urllib.request  
from torchvision.io.image import read_image  

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.

Warning

If the RBLN backend fails to compile the graph, torch.compile() falls back to eager execution on the CPU and only logs a warning. Add 'mode': 'strict' to options to raise an error instead.

 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)

Example Output: 

1
Top-1 category: tabby

References