Arm方案 基于Arm架构的边缘侧设备（树莓派或 NVIDIA Jetson Nano）上部署PyTorch模型

 

               

人工智能 (AI) 在边缘计算领域的应用正迅速普及。因此，在基于 Arm 架构的边缘侧设备上部署机器学习模型变得日益重要。基于 Arm 架构的处理器凭借低功耗和高能效等优势，在嵌入式系统中得到了广泛应用。

本文将为你展示如何在树莓派或 NVIDIA Jetson Nano 等基于 Arm 架构的边缘侧设备上部署 PyTorch 模型。

前提条件

 

在开始之前，请确保准备好以下内容：

 

硬件：一台基于 Arm 架构的设备，例如树莓派、NVIDIA Jetson Nano 或其他类似的边缘侧设备。

 

软件

 

设备上必须安装 Python 3.7 或更高版本。

 

一个与 Arm 架构兼容的 PyTorch 版本。

 

一个经过训练的 PyTorch 模型。

 

依赖项：必须安装诸如 torch 和 torchvision 等库以及其他所需的 Python 包。

第 1 步

准备 PyTorch 模型

 

训练或加载模型

 

在开发机器上训练模型，或从 PyTorch 模型库加载预训练模型：

import torch

import torchvision.models as models

 

# Load a pre-trained model

model = models.resnet18(pretrained=True)

model.eval()

 

优化模型

 

将模型转换为 TorchScript 格式，以获得更好的兼容性和性能：

scripted_model = torch.jit.script(model)

 

torch.jit.save(scripted_model, "resnet18_scripted.pt")

第 2 步

设置基于 Arm 架构的边缘侧设备

 

安装依赖项

 

确保 Arm 设备上已安装 Python。

 

安装 PyTorch

 

使用专门为 Arm 设备构建的版本。例如，树莓派用户可以运行以下命令：

pip install torch torchvision

 

验证安装

import torch

 

print(torch.__version__)

 

print(torch.cuda.is_available()) # Check if CUDA is supported (for devices like Jetson Nano)

第 3 步

将模型部署到设备上

 

传输脚本模型

 

使用 scp 或 USB 驱动器，将模型文件 (resnet18_scripted.pt) 复制到 Arm 设备：

scpresnet18_scripted.pt user@device_ip:/path/to/destination

 

运行推理

 

编写 Python 脚本以加载模型并运行推理：

 import torch

from PIL import Image

from torchvision import transforms

 

# Load the model

model = torch.jit.load("resnet18_scripted.pt")

model.eval()

 

# Preprocess an input image

preprocess = transforms.Compose([

    transforms.Resize(256),

    transforms.CenterCrop(224),

    transforms.ToTensor(),

    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),

])

 

img = Image.open("test_image.jpg")

img_tensor = preprocess(img).unsqueeze(0)  # Add batch dimension

 

# Perform inference

with torch.no_grad():

    output = model(img_tensor)

print("Predicted class:", output.argmax(1).item())

 

第 4 步

针对边缘侧性能进行优化

 

量化

 

使用 PyTorch 的量化技术来减小模型大小并提高推理速度：

from torch.quantization import quantize_dynamic

 

 

 

quantized_model = quantize_dynamic(

 

    model, {torch.nn.Linear}, dtype=torch.qint8

 

)

 

torch.jit.save(quantized_model, "resnet18_quantized.pt")

 

利用硬件加速

 

对于配备 GPU 的设备（如 NVIDIA Jetson Nano），确保使用 CUDA 进行加速计算。

 

安装支持 GPU 的相应 PyTorch 版本。

 

性能基准测试

 

测量延迟和吞吐量，以验证模型在边缘侧设备上的性能：

import time

 

 

 

start_time = time.time()

 

with torch.no_grad():

 

    for _ in range(100):

 

        output = model(img_tensor)

 

end_time = time.time()

 

 

 

print("Average Inference Time:", (end_time - start_time) / 100)

 

第 5 步

大规模部署

 

容器化应用

 

使用 Docker 创建可移植的部署环境。

示例 Dockerfile：

FROM python:3.8-slim

 

 

 

RUN pip install torch torchvision pillow

 

COPY resnet18_scripted.pt /app/

 

COPY app.py /app/

 

WORKDIR /app

 

 

 

CMD ["python", "app.py"]

   

监控与更新

 

实施日志记录和监控，确保应用顺利运行。

 

使用 Prometheus 或 Grafana 等工具获取实时洞察。

结论

 

要在基于 Arm 架构的边缘侧设备上部署 PyTorch 模型，需要对模型进行优化、准备好相应软件并使用合适的硬件。上述步骤可帮助开发者在边缘侧部署 AI 应用，从而在靠近数据生成的位置实现快速、高效的推理。快来动手试试吧！