如何在AWS SageMaker上训练部署YOLOv5并构建完整工作流水线

前置准备 #

• 开通SageMaker、S3、ECR的对应操作权限，获取SageMaker执行角色ARN
• 下载YOLOv5官方代码包到本地
• 将标注完成的YOLO格式数据集上传到S3存储桶的指定路径，包含dataset.yaml配置文件

1. 训练阶段Python入口脚本编写 #

训练入口脚本需要兼容SageMaker的环境参数规范，自动读取挂载的数据集路径、输出训练好的模型到指定目录，参考代码如下：

import argparse
import os
from yolov5 import train

def parse_args():
    parser = argparse.ArgumentParser()
    # 读取SageMaker默认挂载的训练数据路径
    parser.add_argument("--data-path", type=str, default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--weights", type=str, default="yolov5s.pt")
    # 读取SageMaker指定的模型输出路径
    parser.add_argument("--model-dir", type=str, default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()
    # 调用YOLOv5原生训练接口
    train.run(
        data=os.path.join(args.data_path, "dataset.yaml"),
        epochs=args.epochs,
        batch_size=args.batch_size,
        weights=args.weights,
        project=args.model_dir,
        name="train_output",
        exist_ok=True
    )
    # 将最优权重移动到模型目录根目录，方便后续部署读取
    os.system(f"cp {args.model_dir}/train_output/weights/best.pt {args.model_dir}/best.pt")

注意：该脚本需要放在YOLOv5代码根目录下，确保可以正常引用YOLOv5的内部依赖模块

2. SageMaker训练任务启动配置 #

使用SageMaker Python SDK配置并启动训练任务，参考代码如下：

import sagemaker
from sagemaker.pytorch import PyTorch

# 初始化SageMaker会话
sagemaker_session = sagemaker.Session()
role = sagemaker.get_execution_role()

# 定义PyTorch训练器
yolo_estimator = PyTorch(
    entry_point="train.py", # 上一步编写的训练入口脚本名
    source_dir="./yolov5", # 本地YOLOv5代码根目录路径
    role=role,
    instance_count=1,
    instance_type="ml.g4dn.2xlarge", # 可根据训练需求调整GPU实例规格
    framework_version="2.0.0",
    py_version="py310",
    hyperparameters={
        "epochs": 30,
        "batch-size": 16,
        "weights": "yolov5s.pt"
    }
)

# 传入S3数据集路径，启动训练
yolo_estimator.fit({"train": "s3://<你的S3存储桶名>/<数据集目录路径>"})

训练完成后，模型产物会自动上传到S3的对应路径，可通过yolo_estimator.model_data获取路径。

3. 部署阶段Python推理入口脚本编写 #

部署入口脚本需要符合SageMaker的推理接口规范，实现模型加载、输入预处理、推理、输出后处理四个核心函数，参考代码如下：

import torch
import os
import numpy as np
from PIL import Image
from io import BytesIO
from yolov5.models.common import DetectMultiBackend
from yolov5.utils.general import non_max_suppression
from yolov5.utils.augmentations import letterbox

# SageMaker部署时自动调用，用于加载训练好的模型
def model_fn(model_dir):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = DetectMultiBackend(os.path.join(model_dir, "best.pt"), device=device)
    model.eval()
    return model

# 处理输入请求，将传入的图片转为YOLOv5要求的输入格式
def input_fn(request_body, request_content_type):
    if request_content_type == "application/x-image":
        img = Image.open(BytesIO(request_body))
        img = letterbox(np.array(img), 640, stride=32, auto=True)[0]
        img = img.transpose((2, 0, 1))[::-1]
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
        img = img.float() / 255
        if len(img.shape) == 3:
            img = img[None]
        return img
    raise ValueError(f"不支持的Content-Type: {request_content_type}")

# 执行推理
def predict_fn(input_data, model):
    with torch.no_grad():
        pred = model(input_data)
    pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45)
    return pred

# 处理推理结果，转为可序列化的JSON格式返回
def output_fn(prediction, accept):
    if accept == "application/json":
        res = []
        for det in prediction:
            if len(det):
                for *xyxy, conf, cls in reversed(det):
                    res.append({
                        "bbox": [float(i) for i in xyxy],
                        "confidence": float(conf),
                        "class_id": int(cls)
                    })
        return {"detections": res}, accept
    raise ValueError(f"不支持的Accept类型: {accept}")

注意：该脚本同样放在YOLOv5代码根目录下

4. SageMaker模型部署 #

使用训练得到的模型产物，部署为实时推理端点，参考代码如下：

from sagemaker.pytorch import PyTorchModel

# 创建PyTorch模型实例
yolo_model = PyTorchModel(
    model_data=yolo_estimator.model_data, # 训练输出的S3模型路径
    role=role,
    entry_point="inference.py", # 上一步编写的推理入口脚本名
    source_dir="./yolov5",
    framework_version="2.0.0",
    py_version="py310"
)

# 部署为实时端点
predictor = yolo_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g4dn.xlarge", # 可根据推理需求调整为CPU/GPU实例
    endpoint_name="yolov5-object-detection-endpoint"
)

流水线验证 #

部署完成后，可直接传入本地图片调用端点，检查返回的检测结果是否符合预期，确认整个流水线正常运行。

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