跳到主要內容

啟用 GPU 的 TensorFlow 建置於 conda-forge

·4 分鐘閱讀
Wolf Vollprecht
Wolf Vollprecht
conda-forge/core 成員

Tensorflow on Anvil

最近我們已經能夠將啟用 GPU 的 TensorFlow 建置新增到 conda-forge！這是一段漫長的旅程，有多位貢獻者嘗試不同的方法來說服 TensorFlow 基於 Bazel 的建置系統建置啟用 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 推動 cudatoolkit 和 cudnn 在 conda-forge 上實現這一切
  • OVH 慷慨贊助了大型建置機器，我們可以利用這些機器來建置配方
  • Bloomberg 贊助 QuantStack 參與 conda-forge
  • Andreas Trawoger (@atrawog) 提供了此基礎的 Ansible 腳本
  • Thorsten Beier (@derthorsten) 和 Adrien Delsalle (@adriendelsalle) 對配方的貢獻