Dockerfile install `tensorrt-cu12==10.1.0` (#14240)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
pull/14248/head
Glenn Jocher 5 months ago committed by GitHub
parent 209f90b35c
commit 19a2d0a4b0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 3
      docker/Dockerfile
  2. 3
      ultralytics/engine/exporter.py
  3. 3
      ultralytics/nn/autobackend.py

@ -35,7 +35,8 @@ ADD https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt $A
# Install pip packages # Install pip packages
RUN python3 -m pip install --upgrade pip wheel 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 # Run exports to AutoInstall packages
# Edge TPU export fails the first time so is run twice here # Edge TPU export fails the first time so is run twice here

@ -688,7 +688,8 @@ class Exporter:
if LINUX: if LINUX:
check_requirements("tensorrt>7.0.0,<=10.1.0") check_requirements("tensorrt>7.0.0,<=10.1.0")
import tensorrt as trt # noqa 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 # Setup and checks
LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...") LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...")

@ -227,7 +227,8 @@ class AutoBackend(nn.Module):
if LINUX: if LINUX:
check_requirements("tensorrt>7.0.0,<=10.1.0") check_requirements("tensorrt>7.0.0,<=10.1.0")
import tensorrt as trt # noqa 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": if device.type == "cpu":
device = torch.device("cuda:0") device = torch.device("cuda:0")
Binding = namedtuple("Binding", ("name", "dtype", "shape", "data", "ptr")) Binding = namedtuple("Binding", ("name", "dtype", "shape", "data", "ptr"))

Loading…
Cancel
Save