Upgrade PyTorch to 1.2, TorchVison to 0.4 (NVIDIA#1155)

- upgrades PyTorch versions for CUDA 10 - CUDA 9.0 is deprecated in pytroch in favor of the CUDA 9.2 so separate PyTorch and TorchVision versions between CUDA versions Signed-off-by: Janusz Lisiecki <[email protected]>
qmy612 · Aug 9, 2019 · 1cab6f6 · 1cab6f6
1 parent 31051d0
commit 1cab6f6
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/qa/TL0_FW_iterators/test.sh b/qa/TL0_FW_iterators/test.sh
@@ -2,7 +2,7 @@
 # used pip packages
 
 
-pip_packages="nose numpy opencv-python tensorflow-gpu torchvision mxnet-cu##CUDA_VERSION##"
+pip_packages="nose numpy opencv-python tensorflow-gpu torch torchvision mxnet-cu##CUDA_VERSION##"
 target_dir=./dali/test/python
 
 one_config_only=true

diff --git a/qa/TL1_tensorflow_plugin/test.sh b/qa/TL1_tensorflow_plugin/test.sh
@@ -5,8 +5,11 @@ target_dir=./dali/test/python
 
 do_once() {
     USE_CUDA_VERSION=$(cat /usr/local/cuda/version.txt | sed 's/.*Version \([0-9]\+\)\.\([0-9]\+\).*/\1/')
-    test "$USE_CUDA_VERSION" = "10" && export TENSORFLOW_VERSIONS="1.13.1 1.14"
-    test "$USE_CUDA_VERSION" = "9" && export TENSORFLOW_VERSIONS="1.7 1.8 1.9 1.10 1.11 1.12"
+    if [[ "$USE_CUDA_VERSION" = "10" ]]; then
+        export TENSORFLOW_VERSIONS="1.13.1 1.14"
+    else
+        export TENSORFLOW_VERSIONS="1.7 1.8 1.9 1.10 1.11 1.12"
+    fi
 }
 
 test_body() {

diff --git a/qa/setup_packages.py b/qa/setup_packages.py
@@ -24,8 +24,10 @@
             "mxnet-cu90" : ["1.4.0"],
             "mxnet-cu100" : ["1.4.0"],
             "tensorflow-gpu" : {"90": ["1.12.0", "1.11", "1.7"], "100": ["1.13.1", "1.14.0"]},
-            "torch" : ["http://download.pytorch.org/whl/{cuda_v}/torch-1.1.0-{0}.whl"],
-            "torchvision" : ["https://download.pytorch.org/whl/{cuda_v}/torchvision-0.3.0-{0}.whl"]
+            "torch" : {"90": ["http://download.pytorch.org/whl/{cuda_v}/torch-1.1.0-{0}.whl"],
+                       "100": ["http://download.pytorch.org/whl/{cuda_v}/torch-1.2.0-{0}.whl"]},
+            "torchvision" : {"90": ["https://download.pytorch.org/whl/{cuda_v}/torchvision-0.3.0-{0}.whl"],
+                             "100": ["https://download.pytorch.org/whl/{cuda_v}/torchvision-0.4.0-{0}.whl"]},
             }
 
 parser = argparse.ArgumentParser(description='Env setup helper')