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:
| pip3 install torch torchvision
pip install --extra-index-url https://pypi.rbln.ai/simple/ rebel-compiler>=0.8.2
|
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.
| 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.
| 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)
|
Example Output:
Using torch.compile()
API
Prepare the ResNet50 model and input image.
| 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.
| 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:
References