diff --git a/export.py b/export.py index 7ec937c129d5..f0ab22a21670 100644 --- a/export.py +++ b/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 diff --git a/utils/docker/Dockerfile b/utils/docker/Dockerfile index f7554c7b48e4..f59fb02ee12a 100644 --- a/utils/docker/Dockerfile +++ b/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 diff --git a/utils/docker/Dockerfile-cpu b/utils/docker/Dockerfile-cpu index 53cd635aee11..1b613958789e 100644 --- a/utils/docker/Dockerfile-cpu +++ b/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 -------------------------------------------------------------------------------------------------------