Merge pull request #96488 from danieldk/libtorch

libtorch-bin: init at 1.6.0
This commit is contained in:
Daniël de Kok 2020-09-01 13:46:14 +02:00 committed by GitHub
commit 53880bd6c2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 167 additions and 0 deletions

View file

@ -0,0 +1,111 @@
{ callPackage
, stdenv
, fetchzip
, lib
, addOpenGLRunpath
, patchelf
, fixDarwinDylibNames
, cudaSupport
, nvidia_x11
}:
let
version = "1.6.0";
device = if cudaSupport then "cuda" else "cpu";
srcs = import ./binary-hashes.nix;
unavailable = throw "libtorch is not available for this platform";
in stdenv.mkDerivation {
inherit version;
pname = "libtorch";
src = fetchzip srcs."${stdenv.targetPlatform.system}-${device}" or unavailable;
nativeBuildInputs =
if stdenv.isDarwin then [ fixDarwinDylibNames ]
else [ addOpenGLRunpath patchelf ]
++ stdenv.lib.optionals cudaSupport [ addOpenGLRunpath ];
buildInputs = [
stdenv.cc.cc
] ++ lib.optionals cudaSupport [ nvidia_x11 ];
dontBuild = true;
dontConfigure = true;
dontStrip = true;
installPhase = ''
# Copy headers and CMake files.
install -Dm755 -t $dev/lib lib/*.a
cp -r include $dev
cp -r share $dev
install -Dm755 -t $out/lib lib/*${stdenv.hostPlatform.extensions.sharedLibrary}*
# We do not care about Java support...
rm -f $out/lib/lib*jni* 2> /dev/null || true
'';
postFixup = let
libPaths = [ stdenv.cc.cc.lib ]
++ stdenv.lib.optionals cudaSupport [ nvidia_x11 ];
rpath = stdenv.lib.makeLibraryPath libPaths;
in stdenv.lib.optionalString stdenv.isLinux ''
find $out/lib -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
echo "setting rpath for $lib..."
patchelf --set-rpath "${rpath}:$out/lib" "$lib"
${lib.optionalString cudaSupport ''
addOpenGLRunpath "$lib"
''}
done
'' + stdenv.lib.optionalString stdenv.isDarwin ''
install_name_tool -change @rpath/libshm.dylib $out/lib/libshm.dylib $out/lib/libtorch_python.dylib
install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libtorch_python.dylib
install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch_python.dylib
install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libtorch_python.dylib
install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libtorch_python.dylib
install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libtorch.dylib
install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch.dylib
install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libtorch.dylib
install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libtorch_cpu.dylib
install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch_cpu.dylib
install_name_tool -change @rpath/libtensorpipe.dylib $out/lib/libtensorpipe.dylib $out/lib/libtorch_cpu.dylib
install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libcaffe2_observers.dylib
install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libcaffe2_observers.dylib
install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libcaffe2_observers.dylib
install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libcaffe2_observers.dylib
install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libcaffe2_detectron_ops.dylib
install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libcaffe2_detectron_ops.dylib
install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libcaffe2_detectron_ops.dylib
install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libcaffe2_detectron_ops.dylib
install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libshm.dylib
install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libshm.dylib
install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libshm.dylib
install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libshm.dylib
install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch_global_deps.dylib
install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libtorch_global_deps.dylib
'';
outputs = [ "out" "dev" ];
passthru.tests = callPackage ./test { };
meta = with stdenv.lib; {
description = "C++ API of the PyTorch machine learning framework";
homepage = "https://pytorch.org/";
license = licenses.unfree; # Includes CUDA and Intel MKL.
platforms = with platforms; linux ++ darwin;
};
}

View file

@ -0,0 +1,14 @@
{
x86_64-darwin-cpu = {
url = "https://download.pytorch.org/libtorch/cpu/libtorch-macos-1.6.0.zip";
sha256 = "0d4n7la31qzl4s9pwvm07la7q6lhcwiww0yjpfz3kw6nvx84p22r";
};
x86_64-linux-cpu = {
url = "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-1.6.0%2Bcpu.zip";
sha256 = "1975b4zvyihzh89vnwspw0vf9qr05sxj8939vcrlmv3gzvdspcxz";
};
x86_64-linux-cuda = {
url = "https://download.pytorch.org/libtorch/cu102/libtorch-cxx11-abi-shared-with-deps-1.6.0.zip";
sha256 = "127qnfyi1faqbm40sbnsyqxjhrqj82bzwqyz7c1hs2bm0zgrrpya";
};
}

View file

@ -0,0 +1,4 @@
find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
add_executable(test test.cpp)
target_link_libraries(test "${TORCH_LIBRARIES}")

View file

@ -0,0 +1,26 @@
{ stdenv, cmake, libtorch-bin, symlinkJoin }:
stdenv.mkDerivation {
pname = "libtorch-test";
version = libtorch-bin.version;
src = ./.;
postPatch = ''
cat CMakeLists.txt
'';
makeFlags = [ "VERBOSE=1" ];
nativeBuildInputs = [ cmake ];
buildInputs = [ libtorch-bin ];
installPhase = ''
touch $out
'';
checkPhase = ''
./test
'';
}

View file

@ -0,0 +1,7 @@
#include <torch/torch.h>
#include <iostream>
int main() {
torch::Tensor tensor = torch::eye(3);
std::cout << tensor << std::endl;
}

View file

@ -1372,6 +1372,11 @@ in
then python.pkgs.tensorflow.libtensorflow
else libtensorflow-bin;
libtorch-bin = callPackage ../development/libraries/science/math/libtorch/bin.nix {
inherit (linuxPackages) nvidia_x11;
cudaSupport = pkgs.config.cudaSupport or false;
};
behdad-fonts = callPackage ../data/fonts/behdad-fonts { };
bless = callPackage ../applications/editors/bless { };