SageMaker Neo编译问题：无法完成FP16与INT8精度编译

前提准备 #

确保你的PyTorch YOLOv5 Large模型已导出为TorchScript格式（Neo仅支持TorchScript或ONNX作为输入模型）。导出示例：

import torch
from models.yolo import Model
import yaml

# 加载YOLOv5 Large配置和权重
with open('models/yolov5l.yaml', 'r') as f:
    cfg = yaml.safe_load(f)
model = Model(cfg, ch=3, nc=80)
model.load_state_dict(torch.load('yolov5l.pt')['model'].state_dict())
model.eval()

# 固定输入维度导出TorchScript（必须指定固定shape，Neo不支持动态维度）
input_sample = torch.randn(1, 3, 640, 640)
traced_model = torch.jit.trace(model, input_sample)
traced_model.save('yolov5l_traced.pt')

实现FP16精度编译 #

在SageMaker Neo的OutputConfig中，需明确指定precision为fp16，同时确保target_device设置正确。完整的编译参数示例：

import boto3

sm_client = boto3.client('sagemaker')

response = sm_client.create_compilation_job(
    CompilationJobName='yolov5l-fp16-jetson',
    RoleArn='arn:aws:iam::your-account-id:role/SageMakerNeoRole',
    InputConfig={
        'S3Uri': 's3://your-bucket/path/yolov5l_traced.pt',
        'DataInputConfig': '{"input": [1,3,640,640]}',  # 必须匹配模型输入shape
        'Framework': 'PYTORCH'
    },
    OutputConfig={
        'S3OutputLocation': 's3://your-bucket/path/neo-output',
        'TargetDevice': 'jetson_xavier_nx',
        'CompilerOptions': {
            'precision': 'fp16'
        }
    },
    StoppingCondition={
        'MaxRuntimeInSeconds': 3600
    }
)

关键注意点：

• DataInputConfig必须明确固定输入维度，不能留空或动态设置。
• TargetDevice必须精确为jetson_xavier_nx，不能使用模糊值。

实现INT8精度编译 #

INT8编译需要校准数据集来计算量化参数，步骤比FP16复杂：

1. 准备校准数据 #

校准数据需满足：

• 格式：预处理后的numpy数组，shape与模型输入一致（如(N,3,640,640)，N为校准样本数，建议10-50张）。
• 数据类型：float32，需和模型训练时的预处理逻辑一致（如归一化、缩放）。
• 存储：可将校准数据保存为numpy文件（如calibration_data.npy）上传至S3，或直接在代码中传入内存中的numpy数组列表。

2. 配置INT8编译参数 #

在CompilerOptions中指定precision: int8，并传入校准数据路径或数据对象。示例代码：

response = sm_client.create_compilation_job(
    CompilationJobName='yolov5l-int8-jetson',
    RoleArn='arn:aws:iam::your-account-id:role/SageMakerNeoRole',
    InputConfig={
        'S3Uri': 's3://your-bucket/path/yolov5l_traced.pt',
        'DataInputConfig': '{"input": [1,3,640,640]}',
        'Framework': 'PYTORCH'
    },
    OutputConfig={
        'S3OutputLocation': 's3://your-bucket/path/neo-output-int8',
        'TargetDevice': 'jetson_xavier_nx',
        'CompilerOptions': {
            'precision': 'int8',
            'calibration_dataset_s3_uri': 's3://your-bucket/path/calibration_data.npy'  # 校准数据S3路径
        }
    },
    StoppingCondition={
        'MaxRuntimeInSeconds': 3600
    }
)

如果不想上传到S3，也可以直接传入校准数据（适用于小批量样本）：

# 假设calibration_data是形状为(10,3,640,640)的numpy数组
CompilerOptions={
    'precision': 'int8',
    'calibration_data': calibration_data.tolist()
}

关键注意点：

• 校准数据必须覆盖模型常见的输入场景，否则量化后精度损失可能较大。
• 确保预处理逻辑与训练/推理时完全一致，包括图像缩放、归一化、通道顺序等。

常见问题排查 #

• 编译后仍为FP32：检查precision参数是否拼写正确，TargetDevice是否指定为jetson_xavier_nx，DataInputConfig是否明确固定输入维度。
• INT8编译失败：检查校准数据的shape、数据类型是否匹配，样本数量是否足够，S3路径是否有权限访问。
• 模型导出问题：必须使用固定输入维度导出TorchScript，动态shape会导致Neo无法正确量化。

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