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

d7v8vwbk 于 2023-10-23 发布在 Python

关注(0)|答案(1)|浏览(182)

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

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

现在的问题是，我的库可以同时使用tensorflow和tensorflow-gpu，我想提供一个python包，默认情况下，它具有tensorflow作为依赖项，但允许用户选择不同的“风味”（在我的情况下，使用tensorflow-gpu并利用GPU进行Tensor计算）。
我已经浏览了Setuptools的文档，但我找不到任何帮助。作为一个可能的解决方案，我认为使用部分[project.optional-dependencies]，可以解决我的问题，例如设置一个字段如下：

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

我可以运行pip install myLibrary[gpu]来包含可选的依赖项，但是，必需的依赖项（即tensorflow）将永远安装，这不是我所期望的行为。
解决这个问题的另一个激烈的方法是发布两个不同的库（myLibrary和myLibrary-gpu），但我认为这不是一个优雅的管理这种情况的方法，特别是因为库中的代码对于两个版本来说是完全相同的。
任何帮助，将非常感谢！

tensorflow

来源：https://stackoverflow.com/questions/77238991/how-can-i-make-my-python-library-with-tensorflow-dependencies-adaptable-for-both

1条答案

按热度按时间

aoyhnmkz1#

我遇到了同样的问题，很难找到一个优雅的解决方案。这是我目前为止最好的解决方案。将下面的代码复制到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可能不适合每个人。

展开查看全部

赞(0）回复(0）举报 2023-10-23

我来回答

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

1条答案

相关问题

热门标签

最新问答