pytorch: 0.3.1 -> 0.4.0

Signed-off-by: Anders Kaseorg <andersk@mit.edu>
This commit is contained in:
Anders Kaseorg 2018-05-04 18:00:50 -04:00
parent 6db7f92cc2
commit d2f7acf30c
2 changed files with 21 additions and 28 deletions

View File

@ -1,7 +1,7 @@
{ buildPythonPackage,
{ buildPythonPackage, pythonOlder,
cudaSupport ? false, cudatoolkit ? null, cudnn ? null,
fetchFromGitHub, fetchpatch, lib, numpy, pyyaml, cffi, cmake,
git, stdenv, linkFarm, symlinkJoin,
fetchFromGitHub, fetchpatch, lib, numpy, pyyaml, cffi, typing, cmake,
stdenv, linkFarm, symlinkJoin,
utillinux, which }:
assert cudnn == null || cudatoolkit != null;
@ -25,7 +25,7 @@ let
"LD_LIBRARY_PATH=${cudaStub}\${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH} ";
in buildPythonPackage rec {
version = "0.3.1";
version = "0.4.0";
pname = "pytorch";
name = "${pname}-${version}";
@ -34,36 +34,17 @@ in buildPythonPackage rec {
repo = "pytorch";
rev = "v${version}";
fetchSubmodules = true;
sha256 = "1k8fr97v5pf7rni5cr2pi21ixc3pdj3h3lkz28njbjbgkndh7mr3";
sha256 = "12d5vqqaprk0igmih7fwa65ldmaawgijxl58h6dnw660wysc132j";
};
patches = [
(fetchpatch {
# make sure stdatomic.h is included when checking for ATOMIC_INT_LOCK_FREE
# Fixes this test failure:
# RuntimeError: refcounted file mapping not supported on your system at /tmp/nix-build-python3.6-pytorch-0.3.0.drv-0/source/torch/lib/TH/THAllocator.c:525
url = "https://github.com/pytorch/pytorch/commit/502aaf39cf4a878f9e4f849e5f409573aa598aa9.patch";
stripLen = 3;
extraPrefix = "torch/lib/";
sha256 = "1miz4lhy3razjwcmhxqa4xmlcmhm65lqyin1czqczj8g16d3f62f";
})
];
postPatch = ''
substituteInPlace test/run_test.sh --replace \
"INIT_METHOD='file://'\$TEMP_DIR'/shared_init_file' \$PYCMD ./test_distributed.py" \
"echo Skipped for Nix package"
'';
preConfigure = lib.optionalString cudaSupport ''
export CC=${cudatoolkit.cc}/bin/gcc
export CC=${cudatoolkit.cc}/bin/gcc CXX=${cudatoolkit.cc}/bin/g++
'' + lib.optionalString (cudaSupport && cudnn != null) ''
export CUDNN_INCLUDE_DIR=${cudnn}/include
'';
buildInputs = [
cmake
git
numpy.blas
utillinux
which
@ -73,10 +54,10 @@ in buildPythonPackage rec {
cffi
numpy
pyyaml
];
] ++ lib.optional (pythonOlder "3.5") typing;
checkPhase = ''
${cudaStubEnv}${stdenv.shell} test/run_test.sh
${cudaStubEnv}python test/run_test.py --exclude distributed
'';
meta = {

View File

@ -5623,8 +5623,20 @@ in {
};
};
pytorch = callPackage ../development/python-modules/pytorch {
pytorch = let
# Fails with CUDA 9.1 and GCC 6.4:
# https://github.com/pytorch/pytorch/issues/5831
# https://devtalk.nvidia.com/default/topic/1028112
# We should be able to remove this when CUDA 9.2 is released.
cudatoolkit9 = pkgs.cudatoolkit9.override {
gcc6 = pkgs.gcc5;
};
in callPackage ../development/python-modules/pytorch {
cudaSupport = pkgs.config.cudaSupport or false;
cudatoolkit = cudatoolkit9;
cudnn = pkgs.cudnn_cudatoolkit9.override {
inherit cudatoolkit9;
};
};
pytorchWithCuda = self.pytorch.override {