如何在 NixOS 上正确设置带有 GPU 加速的张量流

Question

提前声明，我刚刚开始使用NixOS，所以如果我犯了看似基本的错误，请原谅我。

为了寻找问题的解决方案，我查阅了整个互联网。

到目前为止，这是我的 shell.nix：

with import <nixos> {};

mkShell {
 name = "tensorflow-cuda-shell";
 buildInputs = with python311Packages; [
 cudatoolkit
 cudaPackages.cudnn
 tensorflowWithCuda
];
 shellHook = ''
   export LD_LIBRARY_PATH=${pkgs.stdenv.cc.cc.lib}/lib:${pkgs.cudaPackages.cudatoolkit}/lib:${pkgs.cudaPackages.cudnn}/lib:${pkgs.cudaPackages.cudatoolkit.
lib}/lib:$LD_LIBRARY_PATH
   alias pip="PIP_PREFIX='$(pwd)/_build/pip_packages' TMPDIR='$HOME' \pip"
   export PYTHONPATH="$(pwd)/_build/pip_packages/lib/python3.9/site-packages:$PYTHONPATH"
   export PATH="$(pwd)/_build/pip_packages/bin:$PATH"
   unset SOURCE_DATE_EPOCH
 '';
}

这是 tensorflow 上的 wiki 页面（似乎已经过时）、论坛上的回复以及抓取 GitHub 上的自定义配置的合并。使用这个 shell 脚本，目前需要很长时间才能从源代码编译所有内容。已经至少一个小时了，但它仍然编译二进制文件。有什么办法可以加快速度吗？

任何有关如何通过 GPU 加速运行张量流的想法都深表赞赏。看来这个问题已经很长时间没有被提出了，所有的资源都已经过时了。

Answer 1

天哪，我已经明白了。

我引用了NixOS 讨论中的这篇文章。这只是一个

nix-shell

命令，但它对我有用。现在我想创建

shell.nix

文件。

这是您需要做的：

1. 打开

   ~/.config/nixpkgs/config.nix

   并写入：

{ cudaSupport = true; allowUnfree = true;  }

2. shell.nix

{ pkgs ? import <nixpkgs> {}, config ? {} }:
let
 pythonPackages = pkgs.python3.withPackages (ps: [ ps.tensorflow ]);
in
pkgs.mkShell {
 name = "tf";
 buildInputs = [
   (pythonPackages)
 ];
 shellHook = ''
   export PYTHONPATH="${pythonPackages}:${PYTHONPATH:-}"
 '';
}

（谢谢 ChatGPT！）

3. 现在运行

   nix-shell

4. python

5. 运行：

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

6. ???
7. 利润！