From 6d3d0c5e82aa03dd4c1f11bdbfb9c5f3f3c64bc3 Mon Sep 17 00:00:00 2001
From: pinto0309 <rmsdh122@yahoo.co.jp>
Date: Tue, 13 Apr 2021 01:20:35 +0900
Subject: [PATCH] pytorch_v1.8.1+cuda11.0, cuda11+cudnn8

---
 Dockerfile | 29 +++++++++++++++++++++--------
 README.md  | 50 ++++++++++++++++++++++++++------------------------
 2 files changed, 47 insertions(+), 32 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index c5959e0..1b02937 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM ubuntu:20.04
+FROM nvidia/cuda:11.0.3-cudnn8-devel-ubuntu20.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 ARG OSVER=ubuntu1804
@@ -6,9 +6,12 @@ ARG TENSORFLOWVER=2.4.1
 ARG TENSORRTVER=cuda11.0-trt7.1.3.4-ga-20200617
 ARG OPENVINOVER=2021.3.394
 ARG OPENVINOROOTDIR=/opt/intel/openvino_2021
-ARG TORCHVER=1.7.1+cu110
-ARG TORCHVISIONVER=0.8.2+cu110
-ARG TORCHAUDIOVER=0.7.2
+# PyTorch==1.8.1+cu110
+# TorchVision==0.9.1+cu110
+# TorchAudio==0.8.1
+ARG TORCHVER=1.8.0a0+56b43f4
+ARG TORCHVISIONVER=0.9.0a0+8fb5838
+ARG TORCHAUDIOVER=0.8.0a0+e4e171a
 ARG wkdir=/home/user
 
 # dash -> bash
@@ -49,10 +52,6 @@ RUN pip3 install --upgrade pip \
     && pip install --upgrade onnx-simplifier \
     && pip install --upgrade gdown \
     && pip install --upgrade PyYAML \
-    && pip install torch==${TORCHVER} \
-                        torchvision==${TORCHVISIONVER} \
-                        torchaudio==${TORCHAUDIOVER} \
-                        -f https://download.pytorch.org/whl/torch_stable.html \
     && ldconfig \
     && pip cache purge \
     && apt clean \
@@ -108,6 +107,20 @@ RUN gdown --id 1gsAOLzTxUTMV4vKXKay5z9rutjlmP2BM \
     && apt clean \
     && rm -rf /var/lib/apt/lists/*
 
+# Install Custom PyTorch
+RUN gdown --id 1L257ptjP1EnQCDEHwarrDCZw23n4S8rJ \
+    && pip install torch-${TORCHVER}-cp38-cp38-linux_x86_64.whl \
+    && rm torch-${TORCHVER}-cp38-cp38-linux_x86_64.whl \
+    && gdown --id 1B7dsmZYQdiMDWHuEMz-wCMEb-3yz-H9j \
+    && pip install torchvision-${TORCHVISIONVER}-cp38-cp38-linux_x86_64.whl \
+    && rm torchvision-${TORCHVISIONVER}-cp38-cp38-linux_x86_64.whl \
+    && gdown --id 1Y5ZOkRB0dN8fu9J6jrxSbvAm19wnO2v9 \
+    && pip install torchaudio-${TORCHAUDIOVER}-cp38-cp38-linux_x86_64.whl \
+    && rm torchaudio-${TORCHAUDIOVER}-cp38-cp38-linux_x86_64.whl \
+    && pip cache purge \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/*
+
 # Install Custom TensorFlow (MediaPipe Custom OP, FlexDelegate, XNNPACK enabled)
 RUN gdown --id 1P7MPF_05QUcb_jMkLazy8LxGUpNB9V-b \
     && mv tensorflow-2.4.1-cp38-cp38-linux_x86_64.whl tensorflow-2.4.1-cp38-none-linux_x86_64.whl \
diff --git a/README.md b/README.md
index dabd8fc..b0555c0 100644
--- a/README.md
+++ b/README.md
@@ -10,29 +10,31 @@
 
 ## 2. Model optimization environment to be built
 1. Ubuntu 20.04 x86_64
-2. TensorFlow v2.4.1 (MediaPipe Custom OP, FlexDelegate, XNNPACK enabled)
-3. tflite_runtime v2.4.1 (MediaPipe Custom OP, FlexDelegate, XNNPACK enabled)
-4. edgetpu-compiler
-5. flatc 1.12.0
-6. TensorRT cuda11.0-trt7.1.3.4-ga-20200617
-7. PyTorch 1.7.1+cu110
-8. TorchVision 0.8.2+cu110
-9. TorchAudio 0.7.2
-10. OpenVINO 2021.3.394
-11. tensorflowjs
-12. coremltools
-13. onnx
-14. tf2onnx
-15. tensorflow-datasets
-16. openvino2tensorflow
-17. tflite2tensorflow
-18. onnxruntime
-19. onnx-simplifier
-20. MXNet
-21. gdown
-22. OpenCV 4.5.2-openvino
-23. Intel-Media-SDK
-24. Intel iHD GPU (iGPU) support
+2. CUDA 11.0
+3. cuDNN 8.0
+4. TensorFlow v2.4.1 (MediaPipe Custom OP, FlexDelegate, XNNPACK enabled)
+5. tflite_runtime v2.4.1 (MediaPipe Custom OP, FlexDelegate, XNNPACK enabled)
+6. edgetpu-compiler
+7. flatc 1.12.0
+8. TensorRT cuda11.0-trt7.1.3.4-ga-20200617
+9. PyTorch 1.8.1+cu110
+10. TorchVision 0.9.1+cu110
+11. TorchAudio 0.8.1
+12. OpenVINO 2021.3.394
+13. tensorflowjs
+14. coremltools
+15. onnx
+16. tf2onnx
+17. tensorflow-datasets
+18. openvino2tensorflow
+19. tflite2tensorflow
+20. onnxruntime
+21. onnx-simplifier
+22. MXNet
+23. gdown
+24. OpenCV 4.5.2-openvino
+25. Intel-Media-SDK
+26. Intel iHD GPU (iGPU) support
 
 ## 3. Usage
 ### 3-1. Docker Hub
@@ -48,7 +50,7 @@ $ xhost +local: && \
     -e XDG_RUNTIME_DIR=$XDG_RUNTIME_DIR \
     -e DISPLAY=$DISPLAY \
     --privileged \
-    pinto0309/mtomo:ubuntu2004_tf2.4.1_torch1.7.1_openvino2021.3.394
+    pinto0309/mtomo:ubuntu2004_tf2.4.1_torch1.8.1_openvino2021.3.394
 ```
 
 ### 3-2. Docker Build