FROM nvcr.io/nvidia/pytorch:22.04-py3 (#7680)

* FROM nvcr.io/nvidia/pytorch:22.04-py3

* Update Docker

* Update Docker

* Update Docker

* Update Docker

* Update TRT auto-install

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Cleanup

* Cleanup

* Cleanup cpu

* Cleanup cpu

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

export.py

@@ -218,9 +218,12 @@ def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=F
     # YOLOv5 TensorRT export https://developer.nvidia.com/tensorrt
     try:
         assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e. `python export.py --device 0`'
-        if platform.system() == 'Linux':
-            check_requirements(('nvidia-tensorrt',), cmds=('-U --index-url https://pypi.ngc.nvidia.com',))
-        import tensorrt as trt
+        try:
+            import tensorrt as trt
+        except Exception:
+            if platform.system() == 'Linux':
+                check_requirements(('nvidia-tensorrt',), cmds=('-U --index-url https://pypi.ngc.nvidia.com',))
+            import tensorrt as trt
 
         if trt.__version__[0] == '7':  # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012
             grid = model.model[-1].anchor_grid

utils/docker/Dockerfile

@@ -1,7 +1,11 @@
 # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
 
-# Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
-FROM nvcr.io/nvidia/pytorch:21.10-py3
+# Start FROM NVIDIA PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
+FROM nvcr.io/nvidia/pytorch:22.04-py3
+RUN rm -rf /opt/pytorch  # remove 1.2GB dir
+
+# Downloads to user config dir
+ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Arial.Unicode.ttf /root/.config/Ultralytics/
 
 # Install linux packages
 RUN apt update && apt install -y zip htop screen libgl1-mesa-glx
@@ -9,9 +13,9 @@ RUN apt update && apt install -y zip htop screen libgl1-mesa-glx
 # Install pip packages
 COPY requirements.txt .
 RUN python -m pip install --upgrade pip
-RUN pip uninstall -y torch torchvision torchtext
-RUN pip install --no-cache -r requirements.txt albumentations wandb gsutil notebook \
-    torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
+RUN pip uninstall -y torch torchvision torchtext Pillow
+RUN pip install --no-cache -r requirements.txt albumentations wandb gsutil notebook Pillow>=9.1.0 \
+    torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113
 
 # Create working directory
 RUN mkdir -p /usr/src/app
@@ -21,9 +25,6 @@ WORKDIR /usr/src/app
 COPY . /usr/src/app
 RUN git clone https://github.com/ultralytics/yolov5 /usr/src/yolov5
 
-# Downloads to user config dir
-ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Arial.Unicode.ttf /root/.config/Ultralytics/
-
 # Set environment variables
 ENV OMP_NUM_THREADS=8
 

utils/docker/Dockerfile-cpu

@@ -3,6 +3,9 @@
 # Start FROM Ubuntu image https://hub.docker.com/_/ubuntu
 FROM ubuntu:20.04
 
+# Downloads to user config dir
+ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Arial.Unicode.ttf /root/.config/Ultralytics/
+
 # Install linux packages
 RUN apt update
 RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt install -y tzdata
@@ -24,9 +27,6 @@ WORKDIR /usr/src/app
 COPY . /usr/src/app
 RUN git clone https://github.com/ultralytics/yolov5 /usr/src/yolov5
 
-# Downloads to user config dir
-ADD https://ultralytics.com/assets/Arial.ttf https://ultralytics.com/assets/Arial.Unicode.ttf /root/.config/Ultralytics/
-
 
 # Usage Examples -------------------------------------------------------------------------------------------------------