跳到主要内容

conda-forge 上启用 GPU 的 TensorFlow 构建版本

·4 分钟阅读
Wolf Vollprecht
Wolf Vollprecht
conda-forge/core 成员

Tensorflow on Anvil

最近，我们已经能够将启用 GPU 的 TensorFlow 构建版本添加到 conda-forge！这真是一段旅程，多位贡献者尝试了不同的方法来说服基于 Bazel 的 TensorFlow 构建系统构建启用 CUDA 的软件包。但我们成功了，并且拉取请求已合并。

我们现在有了一个配置，可以为所有 conda-forge 支持的配置（CUDA 10.2、11.0、11.1 和 11.2+）创建启用 CUDA 的 TensorFlow 构建版本。构建 CUDA 软件包需要强大的机器——在一台 32 核机器上，构建单个软件包仍然需要大约 3 个小时。我们的构建矩阵现在包括 12 个启用 CUDA 的软件包和 3 个 CPU 软件包（因为我们需要每个 Python 版本的单独软件包）。正如人们可以想象的那样，这在普通的“家用电脑”上不容易实现。

为此，我们编写了一个 Ansible playbook，让我们能够启动云机器，然后这些机器构建 feedstock（使用 build-locally.py 脚本）。感谢 OVH 的慷慨支持，我们能够同时启动多台 32 核虚拟机来构建不同的 TensorFlow 变体。

我们已经开源了 GitHub 中的 Ansible playbook，并且我们正在努力使其（更）普遍适用于其他长时间运行的构建！

Running 3 builds in parallel on 32 cores ... still takes around 3 hours to finish

随着 TensorFlow 构建版本的就位，conda-forge 现在拥有 PyTorch 和 Tensorflow 这两个最流行的深度学习库的启用 CUDA 的构建版本。

我们仍然缺少 TensorFlow 的 Windows 构建版本（CPU 和 CUDA，不幸的是），并且希望社区能够帮助我们解决这个问题。有一个开放的 PR，但可能需要在 Bazel 中进行一些调整才能使其通过：conda-forge/tensorflow-feedstock#111。

我们希望这些新的 GPU 构建版本将使更多软件包能够添加到 conda-forge 频道！我们已经期待 TensorFlow 的 2.6.2 和 2.7 版本，并在未来添加 Windows 支持。我们希望您喜欢这项工作。

安装​

您现在可以使用 tensorflow-gpu 和 tensorflow-cpu 软件包在启用 GPU（默认）和 CPU 软件包之间进行选择。只需运行

mamba install tensorflow-gpu -c conda-forge
# OR
conda install tensorflow-gpu -c conda-forge

在安装 tensorflow 软件包时，如果本地机器有 GPU，软件包解析现在将默认使用 TensorFlow 的启用 GPU 的构建版本（这些构建版本可以通过版本号开头的“cuda”来识别）。请注意，启用 GPU 的软件包也可以在仅 CPU 的机器上运行，但需要像下面这样覆盖环境变量 CONDA_OVERRIDE_CUDA。如果您当前所在的 HPC 节点（例如登录节点）没有 GPU，但具有 GPU 的计算节点没有互联网访问，则这可能很方便。

CONDA_OVERRIDE_CUDA="11.2" conda install tensorflow cudatoolkit>=11.2 -c conda-forge
# OR
CONDA_OVERRIDE_CUDA="11.2" mamba install tensorflow cudatoolkit>=11.2 -c conda-forge

请注意，您应该选择最适合您的 GPU 的 cudatoolkit 版本；目前，我们有“10.2”、“11.0”、“11.1”和“11.2”构建版本可用，其中“11.2”构建版本与所有 cudatoolkit>=11.2 兼容。您也可以通过像上面那样指定来强制使用特定版本的 cudatoolkit。此外，您可以通过附加软件包名称（如 tensorflow==2.7.0=cuda* 或 tensorflow==2.7.0=cuda112*）来确保获得 TensorFlow 的特定构建版本。如果您想要更精简的“仅 CPU”软件包，则可以直接安装 tensorflow-cpu，或者等效地安装 tensorflow==2.7.0=cpu*。在撰写本文时（2022 年 2 月），在没有 GPU 的机器上，除非像上面那样覆盖，否则始终会获得 -cpu 变体。做出此决定是为了让带宽和资源有限的用户更容易访问。

感谢​

  • 感谢 Mark Harfouche (@hmaarrfk) & Ista Zahn (@izahn) 在 TensorFlow GPU 构建版本上的初步工作，以及所有其他 TensorFlow 维护人员。感谢 Uwe Korn (@xhochy) 在 Bazel 脚本和 TensorFlow 上的工作——以及 TensorFlow feedstock 的所有其他维护人员！
  • 感谢 NVIDIA 在 conda-forge 上推动 cudatoolkit 和 cudnn，使其成为可能
  • 感谢 OVH 慷慨赞助大型构建机器，我们可以使用这些机器来构建 recipe
  • 感谢 Bloomberg 赞助 QuantStack 参与 conda-forge
  • 感谢 Andreas Trawoger (@atrawog) 提供此 Ansible 脚本的基础
  • 感谢 Thorsten Beier (@derthorsten) 和 Adrien Delsalle (@adriendelsalle) 对 recipe 的贡献