Tensorflow

From NixOS Wiki
Revision as of 11:24, 2 May 2024 by Sv-cheats-1 (talk | contribs) (Updated wiki page on tensorflow. Old configuration didn't work anymore, so I provided one that worked for me)

Tensorflow

Quick shell [1]

nix-shell --arg config '{ cudaSupport = true; allowUnfree = true; }' -p 'python3.withPackages (ps: [ ps.tensorflow ])'

This starts a shell with tensorflow and it has GPU support.

Shell config

A basic shell config would look something like this:

# 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:-}"
  '';
}

To make this configuration work, edit ~/.config/nxpkgs/config.nix:

# ~/.config/nxpkgs/config.nix
{ cudaSupport = true; allowUnfree = true;  }

References

[1]