TensorRT Export - "TensorRT: export failure: No module named 'tensorrt'"

[Zud0taki]

Search before asking

• I have searched the YOLOv5 issues and found no similar bug report.

YOLOv5 Component

Export

Bug

I have a problem when trying to export to the engine format.

I have clean-installed YOLOv5 on my machine, as I have done several times before.
Everything seems to work fine for training and detection on my GPU (full PC specs down below).

I have installed the requirements from the export.py as well. Most models can be exported, except the engine one, which I need to work with.
To make that work I have installed CUDA, cudNN and TensorRT as described in the nVidia Tutorials and all checks go through without any errors.
My YOLOv5 venv just doesn´t seem to find the TensorRT installation even though it is connected with CUDA, which is detected by the YOLOv5 venv.

A complete ConsoleLog can be found down below.

Environment

PC Specs:
CPU: AMD Ryzen 9 3900X
GPU: nVidia GeForce RTX 3090
RAM: 64GB DDR-4 @2666MHz

Environment:
Yolo: v6.1-146-gb804b36 torch==1.7.1+cu110 torchvision==0.8.2+cu110
Python: 3.8
CUDA: 11.0
cudNN: 8.4.0
TensorRT: 8.2.1.8

Minimal Reproducible Example

PS C:\Users\KIZwei\Desktop\yolov5> python train.py --img 640 --batch 8 --epochs 3 --data coco128.yaml --weights yolov5s.pt
train: weights=yolov5s.pt, cfg=, data=coco128.yaml, hyp=data\hyps\hyp.scratch-low.yaml, epochs=3, batch_size=8, imgsz=640, rect=False, resume=False, nosave=False, noval=False, noautoanchor=False, evolve=None, bucket=, cache=None, image
_weights=False, device=, multi_scale=False, single_cls=False, optimizer=SGD, sync_bn=False, workers=8, project=runs\train, name=exp, exist_ok=False, quad=False, cos_lr=False, label_smoothing=0.0, patience=100, freeze=[0], save_period=-
1, local_rank=-1, entity=None, upload_dataset=False, bbox_interval=-1, artifact_alias=latest
github: YOLOv5 is out of date by 16 commits. Use git pull or git clone https://github.com/ultralytics/yolov5 to update.
YOLOv5 v6.1-146-gb804b36 torch 1.7.1+cu110 CUDA:0 (NVIDIA GeForce RTX 3090, 24575MiB)

hyperparameters: lr0=0.01, lrf=0.01, momentum=0.937, weight_decay=0.0005, warmup_epochs=3.0, warmup_momentum=0.8, warmup_bias_lr=0.1, box=0.05, cls=0.5, cls_pw=1.0, obj=1.0, obj_pw=1.0, iou_t=0.2, anchor_t=4.0, fl_gamma=0.0, hsv_h=0.01
5, hsv_s=0.7, hsv_v=0.4, degrees=0.0, translate=0.1, scale=0.5, shear=0.0, perspective=0.0, flipud=0.0, fliplr=0.5, mosaic=1.0, mixup=0.0, copy_paste=0.0
Weights & Biases: run 'pip install wandb' to automatically track and visualize YOLOv5 runs (RECOMMENDED)
TensorBoard: Start with 'tensorboard --logdir runs\train', view at http://localhost:6006/

             from  n    params  module                                  arguments

0 -1 1 3520 models.common.Conv [3, 32, 6, 2, 2]
1 -1 1 18560 models.common.Conv [32, 64, 3, 2]
2 -1 1 18816 models.common.C3 [64, 64, 1]
3 -1 1 73984 models.common.Conv [64, 128, 3, 2]
4 -1 2 115712 models.common.C3 [128, 128, 2]
5 -1 1 295424 models.common.Conv [128, 256, 3, 2]
6 -1 3 625152 models.common.C3 [256, 256, 3]
7 -1 1 1180672 models.common.Conv [256, 512, 3, 2]
8 -1 1 1182720 models.common.C3 [512, 512, 1]
9 -1 1 656896 models.common.SPPF [512, 512, 5]
10 -1 1 131584 models.common.Conv [512, 256, 1, 1]
11 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
12 [-1, 6] 1 0 models.common.Concat [1]
13 -1 1 361984 models.common.C3 [512, 256, 1, False]
14 -1 1 33024 models.common.Conv [256, 128, 1, 1]
15 -1 1 0 torch.nn.modules.upsampling.Upsample [None, 2, 'nearest']
16 [-1, 4] 1 0 models.common.Concat [1]
17 -1 1 90880 models.common.C3 [256, 128, 1, False]
18 -1 1 147712 models.common.Conv [128, 128, 3, 2]
19 [-1, 14] 1 0 models.common.Concat [1]
20 -1 1 296448 models.common.C3 [256, 256, 1, False]
21 -1 1 590336 models.common.Conv [256, 256, 3, 2]
22 [-1, 10] 1 0 models.common.Concat [1]
23 -1 1 1182720 models.common.C3 [512, 512, 1, False]
24 [17, 20, 23] 1 229245 models.yolo.Detect [80, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]], [128, 256, 512]]
Model summary: 270 layers, 7235389 parameters, 7235389 gradients, 16.5 GFLOPs

Transferred 349/349 items from yolov5s.pt
Scaled weight_decay = 0.0005
optimizer: SGD with parameter groups 57 weight (no decay), 60 weight, 60 bias
train: Scanning 'C:\Users\KIZwei\Desktop\datasets\coco128\labels\train2017.cache' images and labels... 128 found, 0 missing, 2 empty, 0 corrupt: 100%|██████████| 128/128 [00:00<?, ?it/s]
val: Scanning 'C:\Users\KIZwei\Desktop\datasets\coco128\labels\train2017.cache' images and labels... 128 found, 0 missing, 2 empty, 0 corrupt: 100%|██████████| 128/128 [00:00<?, ?it/s]
Plotting labels to runs\train\exp5\labels.jpg...

AutoAnchor: 4.27 anchors/target, 0.994 Best Possible Recall (BPR). Current anchors are a good fit to dataset
Image sizes 640 train, 640 val
Using 8 dataloader workers
Logging results to runs\train\exp5
Starting training for 3 epochs...

 Epoch   gpu_mem       box       obj       cls    labels  img_size
   0/2      1.8G   0.04504   0.06131   0.02099        75       640: 100%|██████████| 16/16 [00:07<00:00,  2.22it/s]
           Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|██████████| 8/8 [00:01<00:00,  6.85it/s]
             all        128        929      0.824      0.593      0.719      0.467

 Epoch   gpu_mem       box       obj       cls    labels  img_size
   1/2     2.46G   0.04457   0.07212   0.01756       127       640: 100%|██████████| 16/16 [00:01<00:00,  9.38it/s]
           Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|██████████| 8/8 [00:00<00:00, 11.36it/s]
             all        128        929      0.742       0.65      0.734      0.485

 Epoch   gpu_mem       box       obj       cls    labels  img_size
   2/2     2.46G   0.04538   0.06357   0.01874       141       640: 100%|██████████| 16/16 [00:01<00:00,  9.56it/s]
           Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100%|██████████| 8/8 [00:00<00:00, 11.43it/s]
             all        128        929      0.775      0.655      0.744      0.488

3 epochs completed in 0.004 hours.
Optimizer stripped from runs\train\exp5\weights\last.pt, 14.8MB
Optimizer stripped from runs\train\exp5\weights\best.pt, 14.8MB

Validating runs\train\exp5\weights\best.pt...
Fusing layers...
Model summary: 213 layers, 7225885 parameters, 0 gradients, 16.5 GFLOPs
Class Images Labels P R mAP@.5 mAP@.5:.95: 100%|██████████| 8/8 [00:02<00:00, 3.98it/s]
all 128 929 0.772 0.652 0.744 0.488
person 128 254 0.812 0.693 0.795 0.51
bicycle 128 6 1 0.587 0.714 0.422
car 128 46 0.702 0.522 0.582 0.196
motorcycle 128 5 0.681 0.8 0.938 0.608
airplane 128 6 0.942 1 0.995 0.693
bus 128 7 0.61 0.714 0.722 0.641
train 128 3 0.65 0.667 0.753 0.45
truck 128 12 0.691 0.5 0.544 0.264
boat 128 6 0.955 0.333 0.459 0.134
traffic light 128 14 0.617 0.233 0.347 0.213
stop sign 128 2 0.76 1 0.995 0.796
bench 128 9 0.828 0.537 0.65 0.272
bird 128 16 0.902 1 0.995 0.672
cat 128 4 1 0.964 0.995 0.772
dog 128 9 1 0.651 0.816 0.56
horse 128 2 0.815 1 0.995 0.722
elephant 128 17 0.961 0.882 0.924 0.622
bear 128 1 0.669 1 0.995 0.895
zebra 128 4 0.848 1 0.995 0.947
giraffe 128 9 0.718 0.667 0.893 0.682
backpack 128 6 0.749 0.333 0.753 0.358
umbrella 128 18 0.871 0.749 0.9 0.493
handbag 128 19 0.85 0.158 0.32 0.148
tie 128 7 0.854 0.838 0.837 0.542
suitcase 128 4 0.872 1 0.995 0.522
frisbee 128 5 0.655 0.8 0.8 0.697
skis 128 1 0.682 1 0.995 0.398
snowboard 128 7 0.826 0.857 0.885 0.612
sports ball 128 6 0.62 0.667 0.603 0.307
kite 128 10 0.749 0.598 0.611 0.209
baseball bat 128 4 0.405 0.25 0.435 0.189
baseball glove 128 7 0.648 0.429 0.467 0.309
skateboard 128 5 0.712 0.51 0.734 0.477
tennis racket 128 7 1 0.545 0.622 0.287
bottle 128 18 0.588 0.389 0.591 0.312
wine glass 128 16 0.69 0.938 0.904 0.467
cup 128 36 0.776 0.75 0.819 0.506
fork 128 6 0.981 0.333 0.438 0.315
knife 128 16 0.845 0.68 0.715 0.366
spoon 128 22 0.757 0.455 0.619 0.297
bowl 128 28 0.884 0.679 0.779 0.576
banana 128 1 0.823 1 0.995 0.0995
sandwich 128 2 1 0 0.638 0.56
orange 128 4 0.553 1 0.995 0.759
broccoli 128 11 0.35 0.455 0.511 0.369
carrot 128 24 0.699 0.542 0.739 0.499
hot dog 128 2 0.596 1 0.828 0.745
pizza 128 5 0.858 0.8 0.878 0.68
donut 128 14 0.649 1 0.947 0.83
cake 128 4 0.886 1 0.995 0.822
chair 128 35 0.528 0.629 0.603 0.312
couch 128 6 1 0.644 0.891 0.56
potted plant 128 14 0.657 0.714 0.773 0.455
bed 128 3 1 0 0.863 0.593
dining table 128 13 0.858 0.467 0.615 0.401
toilet 128 2 0.813 1 0.995 0.846
tv 128 2 0.626 1 0.995 0.821
laptop 128 3 0.995 0.333 0.665 0.374
mouse 128 2 1 0 0.0717 0.0358
remote 128 8 1 0.621 0.642 0.517
cell phone 128 8 0.668 0.25 0.365 0.191
microwave 128 3 0.837 1 0.995 0.844
oven 128 5 0.391 0.4 0.468 0.311
sink 128 6 0.581 0.248 0.433 0.189
refrigerator 128 5 0.665 0.8 0.769 0.527
book 128 29 0.621 0.31 0.4 0.181
clock 128 9 0.744 0.778 0.888 0.698
vase 128 2 0.54 1 0.995 0.895
scissors 128 1 1 0 0.199 0.0398
teddy bear 128 21 0.866 0.614 0.813 0.506
toothbrush 128 5 0.807 1 0.928 0.562
Results saved to runs\train\exp5
PS C:\Users\KIZwei\Desktop\yolov5> python .\export.py --weights C:\Users\KIZwei\Desktop\yolov5\runs\train\exp5\weights\best.pt --include engine --device 0
export: data=C:\Users\KIZwei\Desktop\yolov5\data\coco128.yaml, weights=['C:\Users\KIZwei\Desktop\yolov5\runs\train\exp5\weights\best.pt'], imgsz=[640, 640], batch_size=1, device=0, half=False, inplace=False, train=False, optim
ize=False, int8=False, dynamic=False, simplify=False, opset=12, verbose=False, workspace=4, nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25, include=['engine']
YOLOv5 v6.1-146-gb804b36 torch 1.7.1+cu110 CUDA:0 (NVIDIA GeForce RTX 3090, 24575MiB)

Fusing layers...
Model summary: 213 layers, 7225885 parameters, 0 gradients, 16.5 GFLOPs

PyTorch: starting from C:\Users\KIZwei\Desktop\yolov5\runs\train\exp5\weights\best.pt with output shape (1, 25200, 85) (14.1 MB)

TensorRT: export failure: No module named 'tensorrt'

Additional

Anything that can resolve this issue is greatly appreciated!
Thank you in advance :)

Are you willing to submit a PR?

• Yes I'd like to help by submitting a PR!

[github-actions]

👋 Hello @Zud0taki, thank you for your interest in YOLOv5 🚀! Please visit our ⭐️ Tutorials to get started, where you can find quickstart guides for simple tasks like Custom Data Training all the way to advanced concepts like Hyperparameter Evolution.

If this is a 🐛 Bug Report, please provide screenshots and minimum viable code to reproduce your issue, otherwise we can not help you.

If this is a custom training ❓ Question, please provide as much information as possible, including dataset images, training logs, screenshots, and a public link to online W&B logging if available.

For business inquiries or professional support requests please visit https://ultralytics.com or email support@ultralytics.com.

Requirements

Python>=3.7.0 with all requirements.txt installed including PyTorch>=1.7. To get started:

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt  # install

Environments

YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Google Colab and Kaggle notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

Status

If this badge is green, all YOLOv5 GitHub Actions Continuous Integration (CI) tests are currently passing. CI tests verify correct operation of YOLOv5 training (train.py), validation (val.py), inference (detect.py) and export (export.py) on macOS, Windows, and Ubuntu every 24 hours and on every commit.

[MiksVasiljevs]

You need to have the TensorRT package for Python, as it is used for the .engine export:

yolov5/export.py

Lines 217 to 223 in 2611477

	def export_engine(model, im, file, train, half, simplify, workspace=4, verbose=False, prefix=colorstr('TensorRT:')):
	# 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 installing within your Yolov5 environment (if using Conda / virtualenv) via:

pip install nvidia-pyindex
pip install nvidia-tensorrt

As outlined here: https://pypi.org/project/nvidia-tensorrt/

[Zud0taki]

@MiksVasiljevs

Thank you for your input!
The first install went through without problems, sadly the second one crashed.

This is the error message:

PS C:\Users\KIZwei\Desktop\yolov5> pip install nvidia-tensorrt
Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com
Collecting nvidia-tensorrt
Downloading nvidia-tensorrt-0.0.1.dev5.tar.gz (7.9 kB)
Preparing metadata (setup.py) ... error
error: subprocess-exited-with-error

× python setup.py egg_info did not run successfully.
│ exit code: 1
╰─> [17 lines of output]
Traceback (most recent call last):
File "", line 2, in
File "", line 34, in
File "C:\Users\KIZwei\AppData\Local\Temp\pip-install-pc_jkjfg\nvidia-tensorrt_830b8a8cac2d4f0dafc72039ee35d13e\setup.py", line 150, in
raise RuntimeError(open("ERROR.txt", "r").read())
RuntimeError:
###########################################################################################
The package you are trying to install is only a placeholder project on PyPI.org repository.
This package is hosted on NVIDIA Python Package Index.

  This package can be installed as:
  ```
  $ pip install nvidia-pyindex
  $ pip install nvidia-tensorrt
  ```
  ###########################################################################################
 
  [end of output]

note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

note: This is an issue with the package mentioned above, not pip.
hint: See above for details.
PS C:\Users\KIZwei\Desktop\yolov5>

If you have any advice on how to resolve this error, please let me know :)

[MiksVasiljevs]

Sorry, I did not realise you are using Windows. The Nvidia TensorRT pip installation instructions do mention that it only supports Linux: https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html#installing-pip

As discussed before (#7571), your best option would be to use the Docker image provided to export to .engine.

[glenn-jocher]

@Zud0taki use Docker image with TRT preinstalled. That's what it's there for.

Environments

YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Google Colab and Kaggle notebooks with free GPU:
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Amazon Deep Learning AMI. See AWS Quickstart Guide
• Docker Image. See Docker Quickstart Guide

[Zud0taki]

@glenn-jocher and @MiksVasiljevs

Thank you for all your support!
I must've overread that point.
I will try this asap and close this thread after verification.

Have a good day :)

[ghost]

I have the same error and I am using linux ubuntu 20.04

[glenn-jocher]

@safaATcurrus update your code, TRT is automatically installed when required