[pytorch] Upgrade PyTorch to 1.13.0 (#2157)

deepjavalibrary · Nov 17, 2022 · 17fc2b4 · 17fc2b4
1 parent 57f0327
commit 17fc2b4
Show file tree

Hide file tree

Showing 9 changed files with 48 additions and 63 deletions.
diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml
@@ -85,8 +85,6 @@ jobs:
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Ppt_version=$PYTORCH_VERSION
           ./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch"
           ./gradlew :engines:pytorch:pytorch-native:cleanJNI
-          ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcu10 -Ppt_version=$PYTORCH_VERSION
-          ./gradlew :engines:pytorch:pytorch-native:cleanJNI
           ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcu11 -Ppt_version=$PYTORCH_VERSION
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v1

diff --git a/.github/workflows/nightly_publish.yml b/.github/workflows/nightly_publish.yml
@@ -154,7 +154,6 @@ jobs:
         run: |
           export PADDLE_VERSION=$(cat gradle.properties | awk -F '=' '/paddlepaddle_version/ {print $2}')
           ./gradlew :integration:test
-          ./gradlew :integration:clean :integration:test -Dai.djl.default_engine=PyTorch
           ./gradlew :integration:clean :integration:test -Dai.djl.default_engine=OnnxRuntime
           ./gradlew :engines:mxnet:mxnet-model-zoo:test
           ./gradlew :engines:pytorch:pytorch-model-zoo:test
@@ -222,9 +221,9 @@ jobs:
       - name: Publish to snapshot repository
         if: ${{ github.event.inputs.mode == '' || github.event.inputs.mode == 'snapshot' }}
         run: |
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.10.0 -Psnapshot
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.11.0 -Psnapshot
-          ./gradlew clean engines:ml:xgboost:publish -Pgpu -Psnapshot 
+          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.12.1 -Psnapshot
+          ./gradlew clean engines:ml:xgboost:publish -Pgpu -Psnapshot
           ./gradlew clean publish -Psnapshot
           cd bom
           ./gradlew publish -Psnapshot
@@ -236,9 +235,9 @@ jobs:
       - name: Publish to staging repository
         if: ${{ github.event.inputs.mode == 'staging' }}
         run: |
-          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.10.0 -P${{ github.event.inputs.mode }}
           ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.11.0 -P${{ github.event.inputs.mode }}
-          ./gradlew clean engines:ml:xgboost:publish -Pgpu -P${{ github.event.inputs.mode }} 
+          ./gradlew clean engines:pytorch:pytorch-jni:publish -Ppt_version=1.12.1 -P${{ github.event.inputs.mode }}
+          ./gradlew clean engines:ml:xgboost:publish -Pgpu -P${{ github.event.inputs.mode }}
           ./gradlew clean publish -P${{ github.event.inputs.mode }}
           cd bom
           ./gradlew publish -P${{ github.event.inputs.mode }}

diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
@@ -29,6 +29,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.StandardCopyOption;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
@@ -119,6 +120,7 @@ private static void loadLibTorch(LibTorch libTorch) {
 
         Set<String> loadLater = new HashSet<>(deferred);
         try (Stream<Path> paths = Files.walk(libDir)) {
+            List<Path> dependants = new ArrayList<>();
             paths.filter(
                             path -> {
                                 String name = path.getFileName().toString();
@@ -128,6 +130,10 @@ private static void loadLibTorch(LibTorch libTorch) {
                                         && name.contains("nvTools")) {
                                     return false;
                                 }
+                                if (name.startsWith("libarm_compute_")) {
+                                    dependants.add(path);
+                                    return false;
+                                }
                                 return !loadLater.contains(name)
                                         && Files.isRegularFile(path)
                                         && !name.endsWith(JNI_LIB_NAME)
@@ -137,6 +143,10 @@ private static void loadLibTorch(LibTorch libTorch) {
                             })
                     .map(Path::toString)
                     .forEach(LibUtils::loadNativeLibrary);
+            for (Path dep : dependants) {
+                loadNativeLibrary(dep.toAbsolutePath().toString());
+            }
+
             if (Files.exists((libDir.resolve("cudnn64_8.dll")))) {
                 loadNativeLibrary(libDir.resolve("cudnn64_8.dll").toString());
                 loadNativeLibrary(libDir.resolve("cudnn_ops_infer64_8.dll").toString());

diff --git a/engines/pytorch/pytorch-jni/build.gradle b/engines/pytorch/pytorch-jni/build.gradle
@@ -20,12 +20,17 @@ processResources {
         String url = "https://publish.djl.ai/pytorch/${ptVersion}/jnilib/${djl_version}"
         def files = [
                 "linux-x86_64/cpu/libdjl_torch.so",
-                "linux-x86_64/cu102/libdjl_torch.so",
                 "linux-x86_64/cpu-precxx11/libdjl_torch.so",
                 "osx-x86_64/cpu/libdjl_torch.dylib",
                 "win-x86_64/cpu/djl_torch.dll"
         ]
-        if (ptVersion.startsWith("1.12.")) {
+        if (ptVersion.startsWith("1.13.")) {
+            files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
+            files.add("linux-x86_64/cu117/libdjl_torch.so")
+            files.add("linux-x86_64/cu117-precxx11/libdjl_torch.so")
+            files.add("win-x86_64/cu117/djl_torch.dll")
+            files.add("osx-aarch64/cpu/libdjl_torch.dylib")
+        } else if (ptVersion.startsWith("1.12.")) {
             files.add("linux-aarch64/cpu-precxx11/libdjl_torch.so")
             files.add("linux-x86_64/cu116/libdjl_torch.so")
             files.add("linux-x86_64/cu116-precxx11/libdjl_torch.so")
@@ -38,10 +43,7 @@ processResources {
             files.add("win-x86_64/cu113/djl_torch.dll")
             files.add("osx-aarch64/cpu/libdjl_torch.dylib")
         } else {
-            files.add("linux-x86_64/cu113/libdjl_torch.so")
-            files.add("linux-x86_64/cu113-precxx11/libdjl_torch.so")
-            files.add("win-x86_64/cu113/djl_torch.dll")
-            files.add("win-x86_64/cu102/djl_torch.dll")
+            throw new GradleException("Unsupported version: ${ptVersion}.")
         }
         String jnilibDir = "${project.projectDir}/jnilib/${djl_version}"
         files.each { entry ->

diff --git a/engines/pytorch/pytorch-native/README.md b/engines/pytorch/pytorch-native/README.md
@@ -39,8 +39,8 @@ Use the following task to build pytorch JNI library for GPU:
 ### Mac/Linux
 
 ```
-# compile CUDA 10.2 version of JNI
-./gradlew compileJNI -Pcu10
+# compile CUDA 11.X version of JNI
+./gradlew compileJNI -Pcu11
 ```
 
 ## Windows

diff --git a/engines/pytorch/pytorch-native/build.cmd b/engines/pytorch/pytorch-native/build.cmd
@@ -6,18 +6,7 @@
 
 set FILEPATH="libtorch"
 set VERSION=%1
-if "%2" == "cpu" (
-    set DOWNLOAD_URL="https://download.pytorch.org/libtorch/cpu/libtorch-win-shared-with-deps-%VERSION%%%2Bcpu.zip"
-) else if "%2" == "cu102" (
-      set DOWNLOAD_URL="https://download.pytorch.org/libtorch/cu102/libtorch-win-shared-with-deps-%VERSION%%%2Bcu102.zip"
-) else if "%2" == "cu111" (
-      set DOWNLOAD_URL="https://download.pytorch.org/libtorch/cu111/libtorch-win-shared-with-deps-%VERSION%%%2Bcu111.zip"
-) else if "%2" == "cu113" (
-      set DOWNLOAD_URL="https://download.pytorch.org/libtorch/cu113/libtorch-win-shared-with-deps-%VERSION%%%2Bcu113.zip"
-) else if "%2" == "cu116" (
-      set DOWNLOAD_URL="https://download.pytorch.org/libtorch/cu116/libtorch-win-shared-with-deps-%VERSION%%%2Bcu116.zip"
-)
-
+set DOWNLOAD_URL="https://download.pytorch.org/libtorch/%2/libtorch-win-shared-with-deps-%VERSION%%%2B%2.zip"
 
 if exist %FILEPATH% (
     echo Found %FILEPATH%
@@ -30,22 +19,11 @@ if exist %FILEPATH% (
     echo Finished downloading libtorch
 )
 
-if "%VERSION%" == "1.12.1" (
-    copy /y src\main\patch\cuda.cmake libtorch\share\cmake\Caffe2\public\
-)
-if "%VERSION%" == "1.11.0" (
-    copy /y src\main\patch\cuda.cmake libtorch\share\cmake\Caffe2\public\
-)
-if "%VERSION%" == "1.10.0" (
-    set PT_OLD_VERSION=1
-)
-if "%VERSION%" == "1.9.1" (
-    set PT_OLD_VERSION=1
-)
+copy /y src\main\patch\cuda.cmake libtorch\share\cmake\Caffe2\public\
 
 if exist build rd /q /s build
 md build\classes
 cd build
 javac -sourcepath ..\..\pytorch-engine\src\main\java\ ..\..\pytorch-engine\src\main\java\ai\djl\pytorch\jni\PyTorchLibrary.java -h include -d classes
-cmake -DCMAKE_PREFIX_PATH=libtorch -DPT_OLD_VERSION=%PT_OLD_VERSION% ..
+cmake -DCMAKE_PREFIX_PATH=libtorch ..
 cmake --build . --config Release
diff --git a/engines/pytorch/pytorch-native/build.gradle b/engines/pytorch/pytorch-native/build.gradle
@@ -15,15 +15,15 @@ boolean isPrecxx11 = project.hasProperty("precxx11")
 boolean isAarch64 = project.hasProperty("aarch64") || System.properties["os.arch"] == "aarch64"
 
 String FLAVOR = "cpu"
-if (project.hasProperty("cu10")) {
-    FLAVOR = "cu102"
-} else if (project.hasProperty("cu11")) {
-    if (VERSION.startsWith("1.8.") || VERSION.startsWith("1.9.")) {
-        FLAVOR = "cu111"
-    } else if (VERSION.startsWith("1.10.") || VERSION.startsWith("1.11.")) {
+if (project.hasProperty("cu11")) {
+    if (VERSION.startsWith("1.11.")) {
         FLAVOR = "cu113"
-    } else {
+    } else if (VERSION.startsWith("1.12.")) {
         FLAVOR = "cu116"
+    } else if (VERSION.startsWith("1.13.")) {
+        FLAVOR = "cu117"
+    } else {
+        throw new GradleException("Unsupported PyTorch version: ${VERSION}")
     }
 }
 String BINARY_ROOT = "${project.buildDir}/download"
@@ -85,19 +85,20 @@ 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
-    if (ver.startsWith("1.8.") || ver.startsWith("1.9.")) {
-        cu11 = "cu111"
-    } else if (ver.startsWith("1.10.") || ver.startsWith("1.11.")) {
+    if (ver.startsWith("1.11.")) {
         cu11 = "cu113"
-    } else {
+    } else if (ver.startsWith("1.12.")) {
         cu11 = "cu116"
+    } else if (ver.startsWith("1.13.")) {
+        cu11 = "cu117"
+    } else {
+        throw new GradleException("Unsupported PyTorch version: ${ver}")
     }
 
     def files = [
             "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",
-            "cu102/libtorch-cxx11-abi-shared-with-deps-${ver}%2Bcu102.zip"    : "cu102/linux-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",
             "cpu/libtorch-shared-with-deps-${ver}%2Bcpu.zip"                  : "cpu-precxx11/linux-x86_64",
@@ -134,7 +135,7 @@ def copyNativeLibToOutputDir(Map<String, String> fileStoreMap, String binaryRoot
         // CPU dependencies
         copy {
             from("${outputDir}/libtorch/lib/") {
-                include "libc10.*", "c10.dll", "libiomp5*.*", "libgomp*.*", "libtorch.*", "libtorch_cpu.*", "torch.dll", "torch_cpu.dll", "fbgemm.dll", "asmjit.dll", "uv.dll"
+                include "libc10.*", "c10.dll", "libiomp5*.*", "libarm_compute*.*", "libgomp*.*", "libtorch.*", "libtorch_cpu.*", "torch.dll", "torch_cpu.dll", "fbgemm.dll", "asmjit.dll", "uv.dll"
             }
             into("${outputDir}/native/lib")
         }
@@ -267,6 +268,7 @@ task uploadS3 {
         exec {
             commandLine "sh", "-c", "find ${BINARY_ROOT} -type f | xargs gzip"
         }
+
         def f = new File("${BINARY_ROOT}/files.txt")
         def uploadDirs = [
                 "${BINARY_ROOT}/cpu/linux-x86_64/native/lib/",
@@ -275,10 +277,9 @@ task 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}/cu102/linux-x86_64/native/lib/",
-                "${BINARY_ROOT}/cu116/linux-x86_64/native/lib/",
-                "${BINARY_ROOT}/cu116/win-x86_64/native/lib/",
-                "${BINARY_ROOT}/cu116-precxx11/linux-x86_64/native/lib/"
+                "${BINARY_ROOT}/cu117/linux-x86_64/native/lib/",
+                "${BINARY_ROOT}/cu117/win-x86_64/native/lib/",
+                "${BINARY_ROOT}/cu117-precxx11/linux-x86_64/native/lib/"
         ]
         uploadDirs.each { item ->
             fileTree(item).files.name.each {

diff --git a/engines/pytorch/pytorch-native/build.sh b/engines/pytorch/pytorch-native/build.sh
@@ -23,7 +23,7 @@ ARCH=$4
 
 if [[ ! -d "libtorch" ]]; then
   if [[ $PLATFORM == 'linux' ]]; then
-    if [[ ! "$FLAVOR" =~ ^(cpu|cu102|cu111|cu113|cu116)$ ]]; then
+    if [[ ! "$FLAVOR" =~ ^(cpu|cu113|cu116|cu117)$ ]]; then
       echo "$FLAVOR is not supported."
       exit 1
     fi
@@ -46,16 +46,13 @@ if [[ ! -d "libtorch" ]]; then
   fi
 fi
 
-if [[ "$VERSION" =~ ^1\.10\..*|^1\.9\..* ]]; then
-  PT_OLD_VERSION=1
-fi
 pushd .
 
 rm -rf build
 mkdir build && cd build
 mkdir classes
 javac -sourcepath ../../pytorch-engine/src/main/java/ ../../pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java -h include -d classes
-cmake -DCMAKE_PREFIX_PATH=libtorch -DPT_OLD_VERSION=${PT_OLD_VERSION} ..
+cmake -DCMAKE_PREFIX_PATH=libtorch ..
 cmake --build . --config Release -- -j "${NUM_PROC}"
 
 if [[ $PLATFORM == 'darwin' ]]; then

diff --git a/gradle.properties b/gradle.properties
@@ -13,7 +13,7 @@ systemProp.org.gradle.internal.publish.checksums.insecure=true
 
 djl_version=0.20.0
 mxnet_version=1.9.1
-pytorch_version=1.12.1
+pytorch_version=1.13.0
 tensorflow_version=2.7.4
 tflite_version=2.6.2
 dlr_version=1.6.0