About pytorch-release

This release includes building script and python wheel packages(.whl) for pytorch and torchvision on aarch64 platform. An example to show how to use the pytroch engine is included as well. Currently it supports the native building on NXP aarch64 platform with BSP SDK.

Build

Get the latest IMX BSP from https://source.codeaurora.org/external/imx/imx-manifest.
Set up the build environment for one of the NXP aarch64 platforms and edit the conf/local.conf to add the following dependency for pytorch native build:

IMAGE_INSTALL_append = " python3-dev python3-pip python3-wheel python3-pillow python3-setuptools python3-numpy python3-pyyaml python3-cffi python3-future cmake ninja packagegroup-core-buildessential git git-perltools libxcrypt libxcrypt-dev"

Build the BSP images like the following command:

$ bitbake imx-image-full

Clone this repo and execute the build script on NXP aarch64 platform to generate wheel packages:

$ cd /path/to/pytorch-release/src
$ git config --global http.sslverify false
$ git config --global user.email "yourname@nxp.com"
$ git config --global user.name "your name"
$ ./build.sh

Installation

$ pip3 install /path/to/torch-1.6.0-cp37-cp37m-linux_aarch64.whl

$ pip3 install /path/to/torchvision-0.7.0-cp37-cp37m-linux_aarch64.whl

API overview

Getting started

Load the model file and take mobilenetv2 model for example:

model = models.MobileNetV2()
model.load_state_dict(torch.load('./mobilenet_v2-b0353104.pth'))
# Put the model in eval mode
model.eval()

Load an image and preprocess it:

img = Image.open(image_file)
img_t = preprocess(img)
batch_t = torch.unsqueeze(img_t, 0)

Perform inference and get the results back:

with torch.no_grad():
    output = model(batch_t)
percentage = torch.nn.functional.softmax(output[0], dim=0) * 100
tensor_sort, index_sort = torch.sort(percentage, descending=True)
[print((classes[idx], percentage[idx].item())) for idx in index_sort[:5]]

Running examples

There is an example located in the examples folder, which require urllib, PIL and maybe some other Python3 modules depending on your image. You may install the missing modules using pip3.

To run the example you may simply execute it using the Python3 interpreter. There are no arguments and the resources are downloaded by the scripts:

$ python3 pytorch_mobilenetv2.py

The output should look similar to the following: