I'm trying to get TensorFlow to work with a NVIDIA GPU on Ubuntu 22.04 LTS by a shell.nix but get caught up on missing library imports upon import tensorflow as tf.

I have the Ubuntu provided NVIDIA Display Driver installed and libnvidia-container works with GPU but I'm trying to nix:ify my life more and more without plunging into NixOS.

Does anyone have a working shell.nix for GPU-accelerated PyTorch/TensorFlow/JAX on a Debian-based distro?

In reply to@carlthome:matrix.org

I'm trying to get TensorFlow to work with a NVIDIA GPU on Ubuntu 22.04 LTS by a shell.nix but get caught up on missing library imports upon import tensorflow as tf.

I have the Ubuntu provided NVIDIA Display Driver installed and libnvidia-container works with GPU but I'm trying to nix:ify my life more and more without plunging into NixOS.

Does anyone have a working shell.nix for GPU-accelerated PyTorch/TensorFlow/JAX on a Debian-based distro?

{
  description = "PyTorch code for deep learning on chest X-rays (CXR)";

  inputs = {
    nixpkgs.url = "github:nixos/nixpkgs/nixos-22.05";
    flake-utils.url = "github:numtide/flake-utils";
  };

  # See https://old.reddit.com/r/Python/comments/npu66t/reproducible_python_environment_with_nix_flake/
  outputs = { self, nixpkgs, flake-utils }: 
    flake-utils.lib.eachSystem [ "x86_64-linux" ] (system:
      let
        python = "python310";
        pkgs = (import nixpkgs {
          inherit system;
          config = {
            allowUnfree = true;
            cudaSupport = true;
          };
          overlays = [
          ];
        });
        # core pkgs are those required to run headless scripts
        corePythonPkgs = ps: with ps; [
          numpy
          pandas
          pytorchWithCuda
          scikit-learn
          scipy
          #torchinfo  # currently does not build on nix 22.05?? as of 2022-09-29
          torchvision
          tqdm
        ];
        corePythonEnv = pkgs.${python}.withPackages corePythonPkgs;
        jupyterPythonEnv = pkgs.${python}.withPackages (ps: with ps;
          ((corePythonPkgs ps) ++ [
            ipympl  # for %matplotlib widget
            jupyterlab
            matplotlib
            # not strictly for jupyter but still useful for development
            black
          ]));
      in rec {
        packages = {
          download = pkgs.stdenv.mkDerivation {
            name = "download";
            propagatedBuildInputs = with pkgs; [
              google-cloud-sdk
              ./download_mimic_cxr_jpg.sh
            ];
          };
        };
        apps = rec {
          default = jupyter;
          jupyter = {
            type = "app";
            # Note that this is not a full command line; do not include
            # arguments. If you would like to provide a command line, see the
            # "foo" example further down in this file
            program = "${jupyterPythonEnv}/bin/jupyter";
          };
          download = with pkgs; {
            type = "app";
            program = "${packages.download}/download_mimic_cxr_jpg.sh";
          };
        };
        devShell = pkgs.mkShell { buildInputs = with pkgs; [ jupyterPythonEnv ]; };
      });
}

so for example. you can run nix develop then

\$ nixGLNvidia python -c "import torch; print(torch.cuda.is\_available())"
True

so for example. you can run nix develop then

$ nixGLNvidia-515.65.01 python -c "import torch; print(torch.cuda.is\_available())"
True

so for example. you can run nix develop then

$ nixGLNvidia-515.65.01 python -c "import torch; print(torch.cuda.is_available())"
True