Tensorflow

Tensorflow

There several possible ways to install tensorflow.

Nix-native packages (recommended)

Use python3Packages.tensorflow or python3Packages.tensorflowWithCuda.

pip install in a nix-shell

Nixpkgs provides multiple versions, however, it is often desirable to be able to install the latest nightly from pip. This can accomplished in the following ways:

• By making a nix-shell

with import <nixpkgs> {};
mkShell {
  name = "tensorflow-cuda-shell";
  buildInputs = with python3.pkgs; [
    pip
    numpy
    setuptools
  ];
  shellHook = ''
    export LD_LIBRARY_PATH=${pkgs.stdenv.cc.cc.lib}/lib:${pkgs.cudatoolkit_10_1}/lib:${pkgs.cudnn_cudatoolkit_10_1}/lib:${pkgs.cudatoolkit_10_1.lib}/lib:$LD_LIBRARY_PATH
    alias pip="PIP_PREFIX='$(pwd)/_build/pip_packages' TMPDIR='$HOME' \pip"
    export PYTHONPATH="$(pwd)/_build/pip_packages/lib/python3.7/site-packages:$PYTHONPATH"
    export PATH="$(pwd)/_build/pip_packages/bin:$PATH"
    unset SOURCE_DATE_EPOCH
  '';
}

Within this shell, pip install tf-nightly should work and provide GPU support. The cuda toolkit version (and the version of Python) can be changed to correspond with the matching tensorflow version.

Note: On NixOS 20.03 and above LD_LIBRARY_PATH no longer contains /run/opengl-driver/lib:/run/opengl-driver-32/lib by default, preventing tensorflow from discovering the Cuda libraries. This can be solved by manually (or using nixGL[1] ) appending your LD_LIBRARY_PATH in the shellHook, or by reverting to the pre-20.03 behavior by setting setLdLibraryPath to true in your hardware opengl configuration under configuration.nix.

hardware.opengl.setLdLibraryPath = true;

See also

• [2]