|
1 | | -FROM ubuntu:20.04 |
| 1 | +FROM nvidia/cuda:11.0.3-cudnn8-devel-ubuntu20.04 |
2 | 2 |
|
3 | 3 | ENV DEBIAN_FRONTEND=noninteractive |
4 | 4 | ARG OSVER=ubuntu1804 |
5 | 5 | ARG TENSORFLOWVER=2.4.1 |
6 | 6 | ARG TENSORRTVER=cuda11.0-trt7.1.3.4-ga-20200617 |
7 | 7 | ARG OPENVINOVER=2021.3.394 |
8 | 8 | ARG OPENVINOROOTDIR=/opt/intel/openvino_2021 |
9 | | -ARG TORCHVER=1.7.1+cu110 |
10 | | -ARG TORCHVISIONVER=0.8.2+cu110 |
11 | | -ARG TORCHAUDIOVER=0.7.2 |
| 9 | +# PyTorch==1.8.1+cu110 |
| 10 | +# TorchVision==0.9.1+cu110 |
| 11 | +# TorchAudio==0.8.1 |
| 12 | +ARG TORCHVER=1.8.0a0+56b43f4 |
| 13 | +ARG TORCHVISIONVER=0.9.0a0+8fb5838 |
| 14 | +ARG TORCHAUDIOVER=0.8.0a0+e4e171a |
12 | 15 | ARG wkdir=/home/user |
13 | 16 |
|
14 | 17 | # dash -> bash |
@@ -49,10 +52,6 @@ RUN pip3 install --upgrade pip \ |
49 | 52 | && pip install --upgrade onnx-simplifier \ |
50 | 53 | && pip install --upgrade gdown \ |
51 | 54 | && pip install --upgrade PyYAML \ |
52 | | - && pip install torch==${TORCHVER} \ |
53 | | - torchvision==${TORCHVISIONVER} \ |
54 | | - torchaudio==${TORCHAUDIOVER} \ |
55 | | - -f https://download.pytorch.org/whl/torch_stable.html \ |
56 | 55 | && ldconfig \ |
57 | 56 | && pip cache purge \ |
58 | 57 | && apt clean \ |
@@ -108,6 +107,20 @@ RUN gdown --id 1gsAOLzTxUTMV4vKXKay5z9rutjlmP2BM \ |
108 | 107 | && apt clean \ |
109 | 108 | && rm -rf /var/lib/apt/lists/* |
110 | 109 |
|
| 110 | +# Install Custom PyTorch |
| 111 | +RUN gdown --id 1L257ptjP1EnQCDEHwarrDCZw23n4S8rJ \ |
| 112 | + && pip install torch-${TORCHVER}-cp38-cp38-linux_x86_64.whl \ |
| 113 | + && rm torch-${TORCHVER}-cp38-cp38-linux_x86_64.whl \ |
| 114 | + && gdown --id 1B7dsmZYQdiMDWHuEMz-wCMEb-3yz-H9j \ |
| 115 | + && pip install torchvision-${TORCHVISIONVER}-cp38-cp38-linux_x86_64.whl \ |
| 116 | + && rm torchvision-${TORCHVISIONVER}-cp38-cp38-linux_x86_64.whl \ |
| 117 | + && gdown --id 1Y5ZOkRB0dN8fu9J6jrxSbvAm19wnO2v9 \ |
| 118 | + && pip install torchaudio-${TORCHAUDIOVER}-cp38-cp38-linux_x86_64.whl \ |
| 119 | + && rm torchaudio-${TORCHAUDIOVER}-cp38-cp38-linux_x86_64.whl \ |
| 120 | + && pip cache purge \ |
| 121 | + && apt clean \ |
| 122 | + && rm -rf /var/lib/apt/lists/* |
| 123 | + |
111 | 124 | # Install Custom TensorFlow (MediaPipe Custom OP, FlexDelegate, XNNPACK enabled) |
112 | 125 | RUN gdown --id 1P7MPF_05QUcb_jMkLazy8LxGUpNB9V-b \ |
113 | 126 | && mv tensorflow-2.4.1-cp38-cp38-linux_x86_64.whl tensorflow-2.4.1-cp38-none-linux_x86_64.whl \ |
|
0 commit comments