如何通过Python调用YOLOv5的train()函数并传入训练参数？

核心思路 #

train()函数依赖parse_opt()生成的argparse Namespace对象作为参数，我们可以通过两种方式把原命令行参数传入：要么给parse_opt()指定参数列表，要么手动构造符合要求的Namespace对象。

方法1：修改train.py直接触发训练 #

找到train.py末尾的主入口代码，替换成直接传入参数的逻辑：

if __name__ == "__main__":
    # 把原命令行参数转换成列表传入parse_opt
    opt = parse_opt([
        "--img", "512",
        "--batch", "14",
        "--epochs", "5000",
        "--data", "neurons.yaml",
        "--weights", "yolov5s.pt",
        "--cache", "ram"
    ])
    train(opt)

运行python train.py即可直接用指定参数启动训练，无需在命令行输入参数。

方法2：在外部脚本调用train()函数 #

新建一个独立脚本（比如run_train.py），导入train.py的相关函数并传入参数：

from train import parse_opt, train

# 用parse_opt解析参数列表，自动补全默认参数
opt = parse_opt([
    "--img", "512",
    "--batch", "14",
    "--epochs", "5000",
    "--data", "neurons.yaml",
    "--weights", "yolov5s.pt",
    "--cache", "ram"
])

# 也可以手动构造Namespace对象（需包含parse_opt定义的所有参数，避免报错）
# import argparse
# opt = argparse.Namespace(
#     img=512,
#     batch=14,
#     epochs=5000,
#     data="neurons.yaml",
#     weights="yolov5s.pt",
#     cache="ram",
#     device=None,  # 保留默认值
#     workers=8,    # 保留默认值
#     # 其他parse_opt中定义的参数按需补充
# )

train(opt)

运行python run_train.py即可触发训练，完全无需使用subprocess启动子进程。

注意事项 #

• 优先用parse_opt([参数列表])的方式，避免手动构造Namespace时遗漏参数导致训练报错。
• 如果需要调整参数，直接修改列表中的对应值即可，和命令行参数的逻辑完全一致。

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