Arm NN 是由 ARM 公司开发的一种可轻松将神经网络模型部署到 ARM Cortex-A CPU 和 Mali GPU 的高性能推理引擎。而 PyTorch 则是一个基于 Python 的深度学习框架,由 Facebook 开发并维护,支持 GPU 和 CPU 计算,广泛应用于研究和应用领域。
在实际应用中,往往需要将 PyTorch 训练好的模型部署到 ARM 架构的设备上进行推理,这时 ARM NN 就可以起到重要的作用。那么,Arm NN 支持 PyTorch 吗?
答案是肯定的。从 Arm NN 19.05 版本开始,已支持 PyTorch 的 ONNX 格式模型转换和推理。具体来说,我们可以将 PyTorch 模型通过 ONNX 格式转化为 ARM NN 支持的 FlatBuffer 格式,然后在 ARM 架构的设备上进行推理,即可实现模型的高效部署。
以下是基于 PyTorch 创建一个简单的神经网络模型,并将其通过 ONNX 转换为 ARM NN 可以识别的 FlatBuffer 格式:
import torch
import torchvision
import torch.onnx as onnx
from PIL import Image
# define the model
model = torchvision.models.resnet18(pretrained=True)
model.eval()
# load the image
img = Image.open("image.jpg")
img = torchvision.transforms.Resize((224, 224))(img)
img = torchvision.transforms.ToTensor()(img)
img = torchvision.transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))(img)
img = img.unsqueeze(0)
# export the model to ONNX format
onnx_model = "resnet18.onnx"
input_names = ["input"]
output_names = ["output"]
torch.onnx.export(model, img, onnx_model, input_names=input_names, output_names=output_names)
# convert ONNX to FlatBuffer format for ARM NN
flatbuffer_model = "resnet18.fbs"
import subprocess
subprocess.call(['onnx2armnn', onnx_model, flatbuffer_model])
在以上代码中,我们首先加载了一个预训练的 ResNet-18 神经网络模型,然后将一张图片进行预处理,最后通过 PyTorch 提供的 onnx.export()