如何通过Poetry实现PyTorch的CPU/GPU双版本依赖配置？

我使用Poetry作为依赖管理工具，开发了一个用于目标检测的小型计算机视觉模型项目，希望将其发布为Python包，让终端用户可以通过pip install my_package[cpu]或pip install my_package[gpu]选择安装CPU或GPU版本。当前我的pyproject.toml配置如下：

[[tool.poetry.source]]
name = "torch"
url = "https://download.pytorch.org/whl/cpu"
priority = "explicit"

[tool.poetry.group.cpu.dependencies]
torch = {version = "^2.1.2+cpu", source = "torch", optional=true}
torchvision = {version = "^0.16.2+cpu", source = "torch", optional=true}

[tool.poetry.group.gpu.dependencies]
torch = {version = "^2.1.2", optional=true}
torchvision = {version = "^0.16.2", optional=true}

[tool.poetry.extras]
cpu = ["torch", "torchvision"]
gpu = ["torch", "torchvision"]

但该配置无法正常工作，请问是否有可行的实现方案？

你的配置核心问题是：extras直接引用torch和torchvision时，Poetry无法区分不同group下的版本/源差异，导致依赖解析冲突。以下是两种可行的修复方案：

方案一：通过Group关联Extras #

修改配置让extras直接绑定对应的依赖组，确保安装指定extra时自动拉取对应环境的依赖：

[[tool.poetry.source]]
name = "torch-cpu"
url = "https://download.pytorch.org/whl/cpu"
priority = "explicit"

[[tool.poetry.source]]
name = "pypi"
url = "https://pypi.org/simple"
priority = "default"

[tool.poetry.group.cpu.dependencies]
torch = {version = "2.1.2+cpu", source = "torch-cpu", optional = true}
torchvision = {version = "0.16.2+cpu", source = "torch-cpu", optional = true}

[tool.poetry.group.gpu.dependencies]
torch = {version = "^2.1.2", source = "pypi", optional = true}
torchvision = {version = "^0.16.2", source = "pypi", optional = true}

[tool.poetry.extras]
cpu = ["my_package[cpu]"]
gpu = ["my_package[gpu]"]

• 把CPU源改名避免与默认PyPI源混淆
• extras通过包名[组名]的写法直接关联对应group，确保依赖解析精准

方案二：给依赖添加命名别名 #

通过别名区分CPU/GPU版本的torch和torchvision，让extras能精准引用：

[[tool.poetry.source]]
name = "torch-cpu"
url = "https://download.pytorch.org/whl/cpu"
priority = "explicit"

[tool.poetry.dependencies]
torch-cpu = { version = "2.1.2+cpu", source = "torch-cpu", optional = true }
torchvision-cpu = { version = "0.16.2+cpu", source = "torch-cpu", optional = true }
torch-gpu = { version = "^2.1.2", optional = true }
torchvision-gpu = { version = "^0.16.2", optional = true }

[tool.poetry.extras]
cpu = ["torch-cpu", "torchvision-cpu"]
gpu = ["torch-gpu", "torchvision-gpu"]

这种方式需要在代码里做兼容导入：

try:
    # 优先尝试GPU版本
    import torch
    import torchvision
except ImportError:
    try:
        # 尝试CPU版本
        import torch_cpu as torch
        import torchvision_cpu as torchvision
    except ImportError:
        raise ImportError("请安装对应版本依赖：pip install my_package[cpu] 或 pip install my_package[gpu]")

额外注意事项 #

• 带后缀的版本号（如2.1.2+cpu）不要加^前缀，语义化版本规则会忽略后缀，可能导致版本匹配异常，直接写固定版本更稳妥
• 发布前用poetry build生成安装包，测试pip install命令是否正常拉取对应依赖
• 在包的README中说明GPU版本需要提前配置对应版本的CUDA环境

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