Classify images with ResNet-18

This example runs a real convolutional network through the toolkit. It takes the standard ResNet-18 weights trained on ImageNet and rebuilds the network from toolkit operations, layer by layer. The result is a classifier for 1,000 classes of objects and animals.

Objectives

• Convert a PyTorch state dict to bfloat16 arrays for the toolkit.
• Build convolution, batch normalisation, pooling and residual blocks from toolkit operations.
• Classify sample images and compare the result with PyTorch in float32.

Before you begin

1. Install the toolkit. The CPU emulator works for this example, so you don’t need hardware.

2. Check that Python can import the toolkit:

   python -c "import qant_native_computing_toolkit as q; print(q.generic.get_available_npus())"

   On the emulator this prints {0: 'cpu_backend'}.

3. Get the examples repository:

   git clone --recurse-submodules https://github.com/Q-ANT-GmbH/qant_native_computing_toolkit_examples.git
   cd qant_native_computing_toolkit_examples/python/image_recognition

4. Install the example’s dependencies:

   pip install -r requirements.txt

   The file pins torch==2.11.0 and torchvision==0.26.0 from the PyTorch CPU index. We ran it with PyTorch 2.14.1.

5. Make sure the machine can reach the internet. The notebook downloads the ResNet-18 weights (44.7 MB) from download.pytorch.org and a sample image from files.qant.com.

Run the example

1. Start Jupyter and open image_recognition_on_npu.ipynb:

   jupyter notebook image_recognition_on_npu.ipynb

2. Run all cells. The notebook classifies a photo of a Maltese dog and prints the top five classes.

To run it without opening a browser, execute the notebook from the command line:

jupyter nbconvert --to notebook --execute image_recognition_on_npu.ipynb --output executed.ipynb

To try the other sample images, edit the cell that sets url. It has commented-out lines for a rabbit and a train.

How it works

Load the weights

The notebook gets the pretrained weights from torchvision and passes them to QantResNet18, which is defined in qant_resnet.py:

w = torchvision.models.ResNet18_Weights.IMAGENET1K_V1.get_state_dict()
qant_model = QantResNet18(w)

The constructor converts each tensor in the state dict to a bfloat16 NumPy array:

w[k] = w_org[k].detach().numpy().astype(ml_dtypes.bfloat16)

Because the toolkit uses the same weight layouts as PyTorch, no reshaping is needed.

Build the network from toolkit operations

qant_resnet.py defines each part of ResNet-18 with toolkit calls. The residual block, called basic_block, starts with a convolution, batch normalisation and ReLU:

out = q_ai.conv_fprop(
    x, w[f"{layer_name}.conv1.weight"], padding=1, stride=stride, dilation=1
)

A second convolution and batch normalisation follow. The block then adds the input back and applies ReLU again:

out += identity
out = q_ai.relu_fprop(out)

The full forward pass uses eight toolkit operations: conv_fprop, batchnorm2d_fprop, relu_fprop, maxpool2d_fprop, adaptive_avgpool2d_fprop, linear_fprop, add_bias_fprop and softmax_fprop. The skip-connection addition (out += identity) is a plain NumPy addition on the host.

Process a batch

QantResNet18.__call__ accepts a PyTorch batch, as torchvision preprocessing produces. It loops over the images, converts each to bfloat16, runs forward, and stacks the results into a float32 tensor. to_class_prob then applies the toolkit softmax to get class probabilities.

Expected output

These are the results on the CPU emulator. We ran the notebook for the dog image, and a separate script for all three sample images. The script also ran stock torchvision ResNet-18 in float32 with the same weights for comparison.

Image	Top class	Toolkit, CPU emulator, bfloat16	PyTorch, float32
Maltese dog	Maltese dog	93.0%	93.3%
Rabbit	hare	99.2%	99.0%
Train	electric locomotive	99.2%	99.2%

For all three images, the order of the top five classes was the same in both implementations. For the dog, the other four classes were Lhasa (3.6%), Shih-Tzu (0.8%), Chihuahua (0.5%) and Yorkshire terrier (0.5%).

Each image took about 7 seconds on the emulator. The emulator computes the matrix products with simple loops on the CPU, so this time says nothing about the speed of a card.

Run the tests

The example includes a test that executes the notebook and checks the top class:

python -m pytest -v test_image_recog.py

The test asserts that the top class is “Maltese dog” with a probability above 0.8. On the emulator it passed in 15.7 seconds.

Troubleshooting

If importing torchvision fails with operator torchvision::nms does not exist, your torch and torchvision versions don’t match. See Troubleshooting.

What’s next

• Add transposed convolutions for an encoder-decoder network in Image segmentation.
• Learn how to spread work across several cards in Parallel execution.