From 19a2d0a4b09f9509629caf89ca6fb8118dc9ba5d Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 6 Jul 2024 12:19:42 +0200 Subject: [PATCH] Dockerfile install `tensorrt-cu12==10.1.0` (#14240) Signed-off-by: Glenn Jocher --- docker/Dockerfile | 3 ++- ultralytics/engine/exporter.py | 3 ++- ultralytics/nn/autobackend.py | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 72f4ff989f..25e9c4e2bd 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -35,7 +35,8 @@ ADD https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt $A # Install pip packages RUN python3 -m pip install --upgrade pip wheel -RUN pip install --no-cache-dir -e ".[export]" "albumentations>=1.4.6" comet pycocotools +# Pin TensorRT-cu12==10.1.0 to avoid 10.2.0 bug https://github.com/ultralytics/ultralytics/pull/14239 (note -cu12 must be used) +RUN pip install --no-cache-dir -e ".[export]" "tensorrt-cu12==10.1.0" "albumentations>=1.4.6" comet pycocotools # Run exports to AutoInstall packages # Edge TPU export fails the first time so is run twice here diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index 3d5a075474..748b70d7a6 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -688,7 +688,8 @@ class Exporter: if LINUX: check_requirements("tensorrt>7.0.0,<=10.1.0") import tensorrt as trt # noqa - check_version(trt.__version__, "7.0.0", hard=True) # require tensorrt>=7.0.0 + check_version(trt.__version__, ">=7.0.0", hard=True) + check_version(trt.__version__, "<=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239") # Setup and checks LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...") diff --git a/ultralytics/nn/autobackend.py b/ultralytics/nn/autobackend.py index d95dcb13b0..fcb80ee6e1 100644 --- a/ultralytics/nn/autobackend.py +++ b/ultralytics/nn/autobackend.py @@ -227,7 +227,8 @@ class AutoBackend(nn.Module): if LINUX: check_requirements("tensorrt>7.0.0,<=10.1.0") import tensorrt as trt # noqa - check_version(trt.__version__, "7.0.0", hard=True) # require tensorrt>=7.0.0 + check_version(trt.__version__, ">=7.0.0", hard=True) + check_version(trt.__version__, "<=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239") if device.type == "cpu": device = torch.device("cuda:0") Binding = namedtuple("Binding", ("name", "dtype", "shape", "data", "ptr"))