Skip to content

Commit

Permalink
[pytorch] Updates PyTorch engine to 2.1.0
Browse files Browse the repository at this point in the history
  • Loading branch information
frankfliu committed Oct 17, 2023
1 parent 2f4ebee commit 740a4c5
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 23 deletions.
8 changes: 7 additions & 1 deletion engines/pytorch/pytorch-jni/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,13 @@ processResources {
"osx-x86_64/cpu/libdjl_torch.dylib",
"win-x86_64/cpu/djl_torch.dll"
]
if (ptVersion.startsWith("2.0.")) {
if (ptVersion.startsWith("2.1.")) {
files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
files.add("linux-x86_64/cu121/libdjl_torch.so")
files.add("linux-x86_64/cu121-precxx11/libdjl_torch.so")
files.add("win-x86_64/cu121/djl_torch.dll")
files.add("osx-aarch64/cpu/libdjl_torch.dylib")
} else if (ptVersion.startsWith("2.0.")) {
files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
files.add("linux-x86_64/cu118/libdjl_torch.so")
files.add("linux-x86_64/cu118-precxx11/libdjl_torch.so")
Expand Down
2 changes: 1 addition & 1 deletion engines/pytorch/pytorch-native/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ add_library(djl_torch SHARED ${SOURCE_FILES})
if(NOT BUILD_ANDROID)
target_link_libraries(djl_torch "${TORCH_LIBRARIES}")
target_include_directories(djl_torch PUBLIC build/include ${JNI_INCLUDE_DIRS} ${UTILS_INCLUDE_DIR})
set_property(TARGET djl_torch PROPERTY CXX_STANDARD 14)
set_property(TARGET djl_torch PROPERTY CXX_STANDARD 17)
# We have to kill the default rpath and use current dir
set(CMAKE_SKIP_RPATH TRUE)
if(${CMAKE_SYSTEM_NAME} MATCHES "Linux")
Expand Down
37 changes: 17 additions & 20 deletions engines/pytorch/pytorch-native/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -88,15 +88,17 @@ def prepareNativeLib(String binaryRoot, String ver) {

def officialPytorchUrl = "https://download.pytorch.org/libtorch"
def aarch64PytorchUrl = "https://djl-ai.s3.amazonaws.com/publish/pytorch"
String cu11
String cuda
if (ver.startsWith("1.11.")) {
cu11 = "cu113"
cuda = "cu113"
} else if (ver.startsWith("1.12.")) {
cu11 = "cu116"
cuda = "cu116"
} else if (ver.startsWith("1.13.")) {
cu11 = "cu117"
cuda = "cu117"
} else if (ver.startsWith("2.0.")) {
cu11 = "cu118"
cuda = "cu118"
} else if (ver.startsWith("2.1.")) {
cuda = "cu121"
} else {
throw new GradleException("Unsupported PyTorch version: ${ver}")
}
Expand All @@ -105,10 +107,10 @@ def prepareNativeLib(String binaryRoot, String ver) {
"cpu/libtorch-cxx11-abi-shared-with-deps-${ver}%2Bcpu.zip" : "cpu/linux-x86_64",
"cpu/libtorch-macos-${ver}.zip" : "cpu/osx-x86_64",
"cpu/libtorch-win-shared-with-deps-${ver}%2Bcpu.zip" : "cpu/win-x86_64",
"${cu11}/libtorch-cxx11-abi-shared-with-deps-${ver}%2B${cu11}.zip": "${cu11}/linux-x86_64",
"${cu11}/libtorch-win-shared-with-deps-${ver}%2B${cu11}.zip" : "${cu11}/win-x86_64",
"${cuda}/libtorch-cxx11-abi-shared-with-deps-${ver}%2B${cuda}.zip": "${cuda}/linux-x86_64",
"${cuda}/libtorch-win-shared-with-deps-${ver}%2B${cuda}.zip" : "${cuda}/win-x86_64",
"cpu/libtorch-shared-with-deps-${ver}%2Bcpu.zip" : "cpu-precxx11/linux-x86_64",
"${cu11}/libtorch-shared-with-deps-${ver}%2B${cu11}.zip" : "${cu11}-precxx11/linux-x86_64"
"${cuda}/libtorch-shared-with-deps-${ver}%2B${cuda}.zip" : "${cuda}-precxx11/linux-x86_64"
]

def aarch64Files = [
Expand Down Expand Up @@ -138,17 +140,12 @@ def copyNativeLibToOutputDir(Map<String, String> fileStoreMap, String binaryRoot
from zipTree(file)
into outputDir
}
// CPU dependencies
copy {
from("${outputDir}/libtorch/lib/") {
include "libc10.*", "c10.dll", "libiomp5*.*", "libarm_compute*.*", "libgomp*.*", "libnvfuser_codegen.so", "libtorch.*", "libtorch_cpu.*", "torch.dll", "torch_cpu.dll", "fbgemm.dll", "asmjit.dll", "uv.dll", "nvfuser_codegen.dll"
}
into("${outputDir}/native/lib")
}
// GPU dependencies
delete "${outputDir}/libtorch/lib/*.lib"
delete "${outputDir}/libtorch/lib/*.a"

copy {
from("${outputDir}/libtorch/lib/") {
include "libtorch_cuda*.so", "torch_cuda*.dll", "libc10_cuda.so", "c10_cuda.dll", "libcaffe2_nvrtc.so", "libnvrtc*.so.*", "libcudart*.*", "*nvToolsExt*.*", "cudnn*.dll", "caffe2_nvrtc.dll", "nvrtc64*.dll", "uv.dll", "libcublas*", "zlibwapi.dll"
include "libarm_compute*", "libc10_cuda.so", "libc10.*", "libcaffe2_nvrtc.so", "libcu*", "libgfortran-*", "libgomp*", "libiomp*", "libnv*", "libopenblasp-*", "libtorch_cpu.*", "libtorch_cuda*.so", "libtorch.*", "asmjit.dll", "c10_cuda.dll", "c10.dll", "caffe2_nvrtc.dll", "cu*.dll", "fbgemm.dll", "nv*.dll", "torch_cpu.dll", "torch_cuda*.dll", "torch.dll", "uv.dll", "zlibwapi.dll"
}
into("${outputDir}/native/lib")
}
Expand Down Expand Up @@ -287,9 +284,9 @@ tasks.register('uploadS3') {
"${BINARY_ROOT}/cpu/win-x86_64/native/lib/",
"${BINARY_ROOT}/cpu-precxx11/linux-aarch64/native/lib/",
"${BINARY_ROOT}/cpu-precxx11/linux-x86_64/native/lib/",
"${BINARY_ROOT}/cu118/linux-x86_64/native/lib/",
"${BINARY_ROOT}/cu118/win-x86_64/native/lib/",
"${BINARY_ROOT}/cu118-precxx11/linux-x86_64/native/lib/"
"${BINARY_ROOT}/cu121/linux-x86_64/native/lib/",
"${BINARY_ROOT}/cu121/win-x86_64/native/lib/",
"${BINARY_ROOT}/cu121-precxx11/linux-x86_64/native/lib/"
]
uploadDirs.each { item ->
fileTree(item).files.name.each {
Expand Down
2 changes: 1 addition & 1 deletion gradle.properties
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ systemProp.org.gradle.internal.publish.checksums.insecure=true

djl_version=0.24.0
mxnet_version=1.9.1
pytorch_version=2.0.1
pytorch_version=2.1.0
tensorflow_version=2.10.1
tflite_version=2.6.2
trt_version=8.4.1
Expand Down

0 comments on commit 740a4c5

Please sign in to comment.