HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

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

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-07-06 12:19:42 +02:00 • committed by GitHub
parent 209f90b35c
commit 19a2d0a4b0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 6 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

3
docker/Dockerfile
View file

@ -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

3
ultralytics/engine/exporter.py
View file

@ -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__}...")

3
ultralytics/nn/autobackend.py
View file

@ -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"))