如何使用Setuptools使具有TensorFlow依赖项的Python库适用于CPU和GPU版本？

Question

我正在构建一个 python 库，它将作为 python 包部署在 PyPI 上。 我用作构建系统的工具是Setuptools。该库应该作为依赖项与 TensorFlow 一起使用，因此将其添加到项目中非常简单：

[project]
...
dependencies = [
    "tensorflow >= 2.11.0",
]

现在，问题是我的库能够同时使用

tensorflow

和

tensorflow-gpu

，我想提供一个 python 包，默认情况下，它具有

tensorflow

作为依赖项，但让用户选择不同的“风味”（在我的例子中，使用

tensorflow-gpu

并利用 GPU 进行张量计算）。

我已经浏览了Setuptools的文档，但找不到任何帮助。 作为一种可能的解决方案，我认为使用

[project.optional-dependencies]

部分可以解决我的问题，例如设置一个字段如下：

[project.optional-dependencies]
gpu = ["tensorflow-gpu >= 2.11.0"]

我可以运行

pip install myLibrary[gpu]

来包含可选的依赖项，但是，所需的依赖项（即

tensorflow

）将被安装，这不是我所期望的行为。

解决此问题的另一种彻底方法是提供两个不同的库（

myLibrary

和

myLibrary-gpu

），但我认为这不是管理这种情况的优雅方法，特别是因为库内的代码完全相同对于两个版本。

如有任何帮助，我们将不胜感激！

Answer 1

我遇到了同样的问题，很难找到一个优雅的解决方案。这是我迄今为止最好的解决方案。将以下代码复制到

setup.py

并修改它，以便您可以根据子进程输出在依赖项列表中附加正确的 TensorFlow 包。

因为 python3+ 附带了

subprocess

:

import subprocess

# Run the 'nvidia-smi' command
command = "nvidia-smi"

# Python >= 3.7
result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

# Otherwise
# result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)


# Check if the command was successful
if result.returncode == 0:
    # Print the output of the 'nvidia-smi' command
    print("nvidia-smi output:")
    print(result.stdout)
    dependencies = ["tensorflow-gpu >= 2.11.0"]
else:
    # Print any errors that occurred
    print("Error running nvidia-smi:")
    print(result.stderr)
    dependencies = [
        "tensorflow >= 2.11.0",
    ]

注意：这不是通用解决方案，因为由于权限问题，可能并非每个人都可以运行

nvidia-smi

。