Merge pull request #23240 from jyp/tf

tensorflow: depend on cudatoolkit 8 and cudnn 5.1
This commit is contained in:
Frederik Rietdijk 2017-03-07 13:22:07 +01:00 committed by GitHub
commit 70160ff208
2 changed files with 8 additions and 6 deletions

View File

@ -3,8 +3,8 @@
, buildPythonPackage , buildPythonPackage
, isPy35, isPy27 , isPy35, isPy27
, cudaSupport ? false , cudaSupport ? false
, cudatoolkit75 ? null , cudatoolkit ? null
, cudnn5_cudatoolkit75 ? null , cudnn ? null
, gcc49 ? null , gcc49 ? null
, linuxPackages ? null , linuxPackages ? null
, numpy , numpy
@ -16,8 +16,8 @@
, zlib , zlib
}: }:
assert cudaSupport -> cudatoolkit75 != null assert cudaSupport -> cudatoolkit != null
&& cudnn5_cudatoolkit75 != null && cudnn != null
&& gcc49 != null && gcc49 != null
&& linuxPackages != null; && linuxPackages != null;
@ -97,7 +97,7 @@ buildPythonPackage rec {
propagatedBuildInputs = with stdenv.lib; propagatedBuildInputs = with stdenv.lib;
[ numpy six protobuf3_2 swig mock ] [ numpy six protobuf3_2 swig mock ]
++ optionals cudaSupport [ cudatoolkit75 cudnn5_cudatoolkit75 gcc49 ]; ++ optionals cudaSupport [ cudatoolkit cudnn gcc49 ];
# Note that we need to run *after* the fixup phase because the # Note that we need to run *after* the fixup phase because the
# libraries are loaded at runtime. If we run in preFixup then # libraries are loaded at runtime. If we run in preFixup then
@ -105,7 +105,7 @@ buildPythonPackage rec {
postFixup = let postFixup = let
rpath = stdenv.lib.makeLibraryPath rpath = stdenv.lib.makeLibraryPath
(if cudaSupport then (if cudaSupport then
[ gcc49.cc.lib zlib cudatoolkit75 cudnn5_cudatoolkit75 [ gcc49.cc.lib zlib cudatoolkit cudnn
linuxPackages.nvidia_x11 ] linuxPackages.nvidia_x11 ]
else else
[ gcc.cc.lib zlib ] [ gcc.cc.lib zlib ]

View File

@ -31278,6 +31278,8 @@ EOF
tensorflowWithCuda = callPackage ../development/python-modules/tensorflow { tensorflowWithCuda = callPackage ../development/python-modules/tensorflow {
cudaSupport = true; cudaSupport = true;
cudatoolkit = pkgs.cudatoolkit8;
cudnn = pkgs.cudnn51_cudatoolkit80;
}; };
tflearn = buildPythonPackage rec { tflearn = buildPythonPackage rec {