如何基于Hydra创建CLI入口统一管理多任务命令？

方法一：Hydra Compose API + Fire 整合（适配已有项目） #

针对已有独立Hydra任务的场景，可通过Hydra手动加载配置的方式，避开@hydra.main()的自动注入限制，结合Fire实现统一CLI。

1. 重构任务函数 #

移除@hydra.main()装饰器，将任务函数改为直接接收cfg参数的普通函数：

# train.py
def run(cfg):
    # 原训练逻辑保留，比如：
    print(f"开始训练，epochs={cfg.args.epochs}")

2. 编写统一CLI入口 #

创建cli.py，用hydra.compose手动加载对应任务的配置，再调用任务函数：

import hydra
from omegaconf import OmegaConf
from your_project.train import run as train_run
from your_project.infer import run as infer_run

def train(**kwargs):
    # 加载训练任务的默认配置，config_path对应原train.py的配置目录
    with hydra.initialize(version_base=None, config_path="configs/train"):
        # 将CLI传入的参数转为Hydra配置覆盖格式
        overrides = [f"args.{k}={v}" for k, v in kwargs.items()]
        cfg = hydra.compose(config_name="default", overrides=overrides)
        train_run(cfg)

def infer(**kwargs):
    with hydra.initialize(version_base=None, config_path="configs/infer"):
        overrides = [f"args.{k}={v}" for k, v in kwargs.items()]
        cfg = hydra.compose(config_name="default", overrides=overrides)
        infer_run(cfg)

if __name__ == "__main__":
    import fire
    # 支持嵌套命令，比如classify/train
    fire.Fire({
        "train": train,
        "infer": infer,
        "classify": {
            "train": train,
            "infer": infer
        }
    })

3. 配置全局命令 #

在setup.py中添加控制台脚本，让yolov5成为全局命令：

from setuptools import setup

setup(
    name="yolov5",
    version="0.1",
    entry_points={
        "console_scripts": [
            "yolov5 = cli:main",
        ]
    },
    # 补充项目依赖、描述等配置
)

完成后即可按需求调用：

yolov5 train --epochs=2
yolov5 classify/infer --conf=0.2

方法二：Hydra原生多命令组（推荐新项目） #

Hydra 1.1+版本支持多命令组，可直接基于原生能力实现统一CLI，无需额外依赖。

1. 配置文件结构 #

创建根配置和命令子配置：

# configs/config.yaml
defaults:
  - _self_
  - command: train  # 默认命令

# configs/command/train.yaml
_target_: your_project.train.run
args:
  epochs: 10
  lr: 0.001

# configs/command/infer.yaml
_target_: your_project.infer.run
args:
  conf: 0.5
  iou: 0.45

2. 主入口实现 #

# main.py
import hydra
from omegaconf import DictConfig

@hydra.main(version_base=None, config_path="configs", config_name="config")
def main(cfg: DictConfig) -> None:
    # 动态实例化并调用对应任务函数
    task_fn = hydra.utils.instantiate(cfg.command)
    task_fn(cfg.command)

if __name__ == "__main__":
    main()

3. 适配简洁调用格式 #

若要实现yolov5 train这类无参数前缀的调用，可通过argparse做一层包装：

# cli.py
import argparse
import sys
from main import main

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("command", choices=["train", "infer"])
    args, remaining_args = parser.parse_known_args()
    # 转换为Hydra可识别的配置参数
    sys.argv = [sys.argv[0]] + [f"command={args.command}"] + remaining_args
    main()

同样在setup.py中配置入口点后，即可使用：

yolov5 train --args.epochs=2
yolov5 infer --args.conf=0.2

关键注意事项 #

• 方法一适合已有多独立Hydra任务的项目，改动小、灵活性高。
• 方法二更贴合Hydra原生设计，适合新项目从零规划。
• 避免给每个@hydra.main()函数单独加入口点，会造成命令冲突和全局CLI空间污染。

内容的提问来源于stack exchange，提问作者Ayush Chaurasia