如何在NixOS上正确配置带GPU加速的TensorFlow？编译过慢求助

一、加速编译的方法 #

• 启用官方二进制缓存：确保/etc/nix/nix.conf中包含以下配置，直接下载预编译包而非本地编译：

  substituters = https://cache.nixos.org/
  trusted-public-keys = cache.nixos.org-1:6NCHdD59X431o0gWypbMrAURkbJ16ZPMQFGspcDShjY=
  

• 提升并行编译能力：
  • 全局配置：在/etc/nix/nix.conf添加max-jobs = auto（自动匹配CPU核心数），或手动设置为核心数+1（如max-jobs = 8）。
  • 临时设置：启动shell时指定并行数，比如NIX_BUILD_CORES=8 nix-shell。
  • 配置文件内设置：在mkShell中添加enableParallelBuilding = true;。
• 使用稳定版本依赖：选择Nixpkgs稳定分支（如23.11）中的tensorflowWithCuda，这类版本的预编译包覆盖率更高，减少编译需求。

二、可行的TensorFlow GPU加速配置方案 #

以下是简化且可靠的shell.nix配置，避免手动处理依赖冲突：

{ pkgs ? import <nixos> {} }:

pkgs.mkShell {
  name = "tensorflow-cuda-shell";
  buildInputs = with pkgs.python311Packages; [
    # 直接使用nixpkgs预整合的tensorflowWithCuda，已包含CUDA、cuDNN依赖
    tensorflowWithCuda
    # 按需添加其他Python依赖
    numpy
  ];

  shellHook = ''
    # 配置本地pip包目录，避免全局污染
    export PIP_PREFIX="$(pwd)/_build/pip_packages"
    export PYTHONPATH="$PIP_PREFIX/lib/python3.11/site-packages:$PYTHONPATH"
    export PATH="$PIP_PREFIX/bin:$PATH"
    unset SOURCE_DATE_EPOCH
  '';
}

配置说明 #

• 无需手动添加cudatoolkit和cudnn：tensorflowWithCuda已默认整合对应版本的CUDA栈依赖，Nix会自动处理库路径，不用手动设置LD_LIBRARY_PATH。
• 修正Python版本匹配：PYTHONPATH路径与使用的python311Packages对应，避免版本不兼容问题。
• 验证GPU可用性：进入shell后运行以下代码确认：

  import tensorflow as tf
  print(tf.config.list_physical_devices('GPU'))
  

  输出GPU设备信息即表示配置成功。

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