如何在Jetson边缘设备运行SageMaker AutoML训练的图像分类模型？

步骤1：从SageMaker导出完整的AutoML模型包 #

SageMaker AutoML训练完成后，会生成包含模型权重、推理代码和依赖配置的模型包，你需要把它下载到本地：

• 登录SageMaker控制台，找到你的AutoML训练任务，进入模型输出的S3存储路径，下载整个模型文件夹（通常包含model.pth/model.pt、inference.py、requirements.txt等文件）。
• 如果模型封装在指定容器内，可先拉取镜像到本地，再从容器中复制模型文件：

  docker run -it --name automl-model-container 167316028190.dkr.ecr.us-east-2.amazonaws.com/sm-automl-robotorch:1 /bin/bash
  # 在容器内的/opt/ml/model路径下找到模型文件，复制到本地
  docker cp automl-model-container:/opt/ml/model ./local-automl-model
  

步骤2：配置Jetson设备的运行环境 #

Jetson基于ARM64架构，需适配对应版本的依赖：

• 先安装JetPack SDK，确保CUDA、cuDNN、TensorRT等加速库配置完成。
• 安装适配Jetson的PyTorch版本（因AutoML镜像基于robotorch，本质为PyTorch衍生框架）：根据JetPack版本从NVIDIA官方获取对应安装包，比如JetPack 5.1对应PyTorch 2.0.0。
• 安装模型依赖：根据导出的requirements.txt安装适配ARM64的包，替换x86专属库为ARM版本：

  pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu118
  

步骤3：移植模型并编写Jetson适配的推理脚本 #

• 将下载的模型文件和推理代码复制到Jetson设备。
• 修改推理脚本适配Jetson硬件加速：
  • 加载模型时切换到CUDA设备：

    import torch
    model = torch.load('model.pth')
    model = model.to('cuda')
    model.eval()
    

  • 可选：转换为TorchScript格式提升推理速度：

    example_input = torch.randn(1, 3, 224, 224).to('cuda') # 匹配你的模型输入尺寸
    traced_model = torch.jit.trace(model, example_input)
    traced_model.save('traced_model.pt')
    # 后续推理使用转换后的模型
    model = torch.jit.load('traced_model.pt').to('cuda').eval()
    

  • 确保图像预处理逻辑和训练时一致（尺寸、归一化参数等），并将张量移至CUDA设备推理。

步骤4：测试模型推理 #

编写测试脚本验证运行效果：

import cv2
import torch
import torchvision.transforms as transforms

# 加载模型
model = torch.jit.load('traced_model.pt').to('cuda').eval()

# 预处理变换（与训练阶段一致）
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 读取并处理测试图像
img = cv2.imread('test_image.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_tensor = transform(img).unsqueeze(0).to('cuda')

# 执行推理
with torch.no_grad():
    outputs = model(img_tensor)
    pred = torch.argmax(outputs, dim=1).item()
print(f"预测类别: {pred}")

额外优化建议 #

• 用torch2trt库将TorchScript模型转换为TensorRT引擎，进一步利用Jetson硬件加速能力。
• 减少张量在CPU和CUDA设备间的拷贝，尽量在CUDA上完成所有计算环节。

内容的提问来源于stack exchange，提问作者Stateful