如何将预训练Yolo的model.pt上传MLFlow并打包Docker推送AWS ECR

一、正确上传预训练YOLO模型到MLFlow生成制品 #

你之前生成的配置包含多余的scikit-learn依赖，是因为误用了默认的模型日志方法，MLFlow会默认添加对应日志方法关联的依赖，按以下步骤操作即可自定义依赖：

• 第一步：先编写符合你YOLO版本运行要求的自定义conda环境配置文件，命名为yolo_conda_env.yaml，示例如下（可根据你的实际依赖版本调整）：

channels:
  - conda-forge
  - pytorch
dependencies:
  - python=3.10
  - pytorch>=2.0.0
  - torchvision>=0.15.0
  - pip
  - pip:
    - mlflow>=2.3.0
    - ultralytics>=8.0.200 # 若使用YOLOv5/v7等其他版本，替换为对应依赖包即可
    - opencv-python>=4.8.0
    - cloudpickle>=2.2.0
name: yolo-mlflow-runtime

• 第二步：使用PyTorch专属的MLFlow日志接口上传模型，避免自动引入无关依赖，代码示例如下：

import mlflow
import mlflow.pytorch
from ultralytics import YOLO

# 填写你的MLFlow追踪服务地址，本地运行可设为./mlruns
mlflow.set_tracking_uri("你的MLFlow追踪地址")

with mlflow.start_run(run_name="yolo_pretrained_upload") as run:
    # 加载本地的model.pt预训练模型
    yolo_model = YOLO("model.pt")
    
    # 上传模型到MLFlow制品库，指定自定义的conda环境配置
    mlflow.pytorch.log_model(
        pytorch_model=yolo_model.model,
        artifact_path="yolo_pretrained_model",
        conda_env="./yolo_conda_env.yaml",
        registered_model_name="yolo_prod_model" # 可选，直接注册到MLFlow模型注册中心
    )
    print(f"模型制品存储路径：{run.info.artifact_uri}/yolo_pretrained_model")

• 第三步：验证生成的制品配置，进入对应MLFlow运行记录的制品页，检查conda.yaml的依赖是否和你自定义的内容一致。

二、打包依赖为Docker镜像推送至AWS ECR #

• 第一步：从MLFlow制品库下载完整的模型制品包，包含模型文件、conda环境配置等，在同目录下编写Dockerfile，示例如下：

# 使用PyTorch官方基础镜像，减少依赖安装成本
FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime

WORKDIR /app

# 安装YOLO运行所需的系统依赖
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgl1-mesa-glx \
    libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

# 复制MLFlow导出的conda环境配置
COPY conda.yaml .

# 安装conda和pip依赖
RUN conda install --yes --file <(grep -vE '(pip:|channels:|dependencies:|name:)' conda.yaml | grep -v '^ *-' | sed 's/^ *//') \
    && pip install --no-cache-dir -r <(grep -A100 'pip:' conda.yaml | grep -v 'pip:' | sed 's/^ *- //')

# 复制模型文件到镜像内
COPY yolo_pretrained_model /app/model

# 可选：复制你的推理服务入口脚本
COPY inference.py /app/inference.py

EXPOSE 8080

# 自定义启动命令
CMD ["python", "inference.py"]

• 第二步：执行本地镜像构建：
  docker build -t yolo-mlflow-service:v1 .
• 第三步：登录你的AWS ECR仓库，替换尖括号内的参数为你自己的AWS配置：
  aws ecr get-login-password --region <AWS区域> | docker login --username AWS --password-stdin <AWS账号ID>.dkr.ecr.<AWS区域>.amazonaws.com
• 第四步：给本地镜像打ECR标签：
  docker tag yolo-mlflow-service:v1 <AWS账号ID>.dkr.ecr.<AWS区域>.amazonaws.com/<你的ECR仓库名称>:yolo-service-v1
• 第五步：推送镜像到ECR：
  docker push <AWS账号ID>.dkr.ecr.<AWS区域>.amazonaws.com/<你的ECR仓库名称>:yolo-service-v1

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