YOLOv5模型注册MLflow失败的排查及存加载预测方案

我的目标 #

• 将YOLOv5模型保存至MLflow Model Registry
• 直接从模型仓库加载模型进行预测
• 获取预测结果（边界框）

已执行步骤 #

• 克隆YOLOv5仓库
• 修改yolov5文件夹下的train.py文件，加入MLflow参数日志代码
• 运行mlflow ui后发现模型未被记录
• 按照相关博客操作后，实验板块成功记录参数与权重，但模型未注册至模型板块
• 在train.py的main函数中添加以下代码后问题仍未解决：

# Train
if not opt.evolve:
    #---------------------------------------------------------------------------(added by me)
    with mlflow.start_run(run_name=opt.name, nested=True):
        _, model = train(opt.hyp, opt, device, callbacks) #changed the return values
    #---------------------------------------------------------------------------

mlflow.pytorch.log_model(
    model,
    'yolov5_cloned_exp',
    registered_model_name='yolo_cloned_github'
)

当前问题 #

仅实验板块记录参数，模型未被注册至MLflow Model Registry

需求 #

1. 分析操作错误之处
2. 提供YOLOv5在MLflow中保存与加载的替代方法
3. 提供相关参考链接

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