人工智能 (AI) 在邊緣計算領域的應用正迅速普及。因此，在基于 Arm 架構的邊緣側設備上部署機器學習模型變得日益重要。基于 Arm 架構的處理器憑借低功耗和高能效等優(yōu)勢，在嵌入式系統(tǒng)中得到了廣泛應用。

本文將為你展示如何在樹莓派或 NVIDIA Jetson Nano 等基于 Arm 架構的邊緣側設備上部署 PyTorch 模型。

前提條件

在開始之前，請確保準備好以下內(nèi)容：

硬件：一臺基于 Arm 架構的設備，例如樹莓派、NVIDIA Jetson Nano 或其他類似的邊緣側設備。

軟件

設備上必須安裝 Python 3.7 或更高版本。

一個與 Arm 架構兼容的 PyTorch 版本。

一個經(jīng)過訓練的 PyTorch 模型。

依賴項：必須安裝諸如 torch 和 torchvision 等庫以及其他所需的 Python 包。

第 1 步

準備 PyTorch 模型

訓練或加載模型

在開發(fā)機器上訓練模型，或從 PyTorch 模型庫加載預訓練模型：

import torch

import torchvision.models as models

# Load a pre-trained model

model = models.resnet18(pretrained=True)

model.eval()

優(yōu)化模型

將模型轉(zhuǎn)換為 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 驅(qū)動器，將模型文件 (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 步

針對邊緣側性能進行優(yōu)化

量化

使用 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 步

大規(guī)模部署

容器化應用

使用 Docker 創(chuàng)建可移植的部署環(huán)境。

示例 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"]

監(jiān)控與更新

實施日志記錄和監(jiān)控，確保應用順利運行。

使用 Prometheus 或 Grafana 等工具獲取實時洞察。

結論

要在基于 Arm 架構的邊緣側設備上部署 PyTorch 模型，需要對模型進行優(yōu)化、準備好相應軟件并使用合適的硬件。上述步驟可幫助開發(fā)者在邊緣側部署 AI 應用，從而在靠近數(shù)據(jù)生成的位置實現(xiàn)快速、高效的推理。快來動手試試吧！