TensorRT not working for custom trained models

[Nit-Rathore]

I had set up TensorRT 8.6.1 on my Ubuntu 22.04 with CUDA 12.2 and cudNN 9.0.0

I used this to speed up my yolo models and as given in the documentation I was able to achieve significant boost for the image of bus given and when I tried the same for a pre-trained model for object detection it worked the same but when I attempt to try it on a custom trained model for the same purpose of object detection it creates the .engine file but it fails to execute it and throws errors.

My code is quite straightforward: I simply read the path of the video file, generate the .engine files, and then attempt to run model.predict(video_path). I've tried various combinations of parameters, including running model() with stream=True or using model.predict(), with or without saving the video but the problem persists. I did check the issues and in one similar issues someone mentioned that specifying the task helps, but it did not in my case.

Based on observation, I feel it is an issue with .export() for custom trained models.

[glenn-jocher]

Hello! Thanks for reaching out and providing detailed information about the issue you're encountering with TensorRT and your custom trained model. 🌟

It seems like you've already tried a few troubleshooting steps, which is great. Given the details, it might indeed be related to the .export() process for custom models. Here are a couple of suggestions:

1. Ensure Compatibility: Double-check that your custom model's architecture is fully compatible with TensorRT. Some custom layers or configurations might require additional handling or might not be supported directly.

2. Verbose Export: Try exporting your model with verbose logging enabled to get more insights into what might be going wrong during the export process. You can do this by adding the verbose=True argument to your .export() call. This might provide more clues on the issue.

3. TensorRT Version: Although it seems you're using a recent version of TensorRT, ensure that it's fully compatible with the CUDA and cuDNN versions you're using. Sometimes, specific combinations can lead to unexpected behavior.

4. Simplify Model: If possible, try simplifying your model or using a more standard architecture as a test case to see if the issue persists. This can help isolate whether the problem is with the custom aspects of your model.

If you continue to face issues, it would be helpful to share the specific error messages you're encountering during the .engine file execution. This can provide more context for further troubleshooting.

Here's a quick example of how you might enable verbose logging during export:

from ultralytics import YOLO

# Load your custom trained model
model = YOLO('path/to/your/custom_model.pt')

# Export with verbose logging
model.export(format='engine', verbose=True)

We're here to help, so feel free to share any additional details or error logs! 🛠

[Nit-Rathore]

Hey I checked the matrix compatibility and the thing is for pre-trained models I am getting the speedup using TensorRT, but not working on custom-trained models.

`import time
from ultralytics import YOLO

model = YOLO('models/tennis_racket_detect.pt')
model.export(format='engine',verbose=True)

tensorrt_engine_path = "models/tennis_racket_detect.engine"
video_path = "test_videos/3.mp4"
tensorrt_model = YOLO(tensorrt_engine_path,task='detect')

start_time = time.time()
results = tensorrt_model.predict(video_path)
end_time = time.time()
print(f"TensorRT Video detection time for '{video_path}': {end_time - start_time:.4f} seconds")
`

#Attaching verbose output and the following error message when trying the model

[Nit-Rathore]

Let me know if there is any other information that helps to debug the issue?

[glenn-jocher]

Hello! Thanks for reaching out and providing detailed information about the issue you're encountering with TensorRT and your custom-trained model. 🚀

From the details you've shared, it seems like the export process to TensorRT format is successful, but you're facing challenges during inference. A common reason for such issues could be related to the input dimensions or the specific configurations used during the training of your custom model that might not align perfectly with the TensorRT optimization process.

To help us dive deeper into this, could you please share the following additional details:

• The exact error message you're receiving during inference.
• The input size (imgsz) you used during the training of your custom model.
• Any specific custom layers or modifications made to the model architecture.

In the meantime, you might want to try exporting your model with different imgsz parameters that match your training settings, if you haven't already. For example:

model.export(format='engine', imgsz=640, verbose=True)

This could sometimes resolve discrepancies between the training and inference phases. Looking forward to your response so we can assist you further!

[Nit-Rathore]

• Image of the exact error message during inference:

• The input size used for training was 640 and we took the pre-trained model from YoloV8 and trained it on dataset from roboflow having about 1000 images for 300 epochs.
  YAML file for the same is below, I referred the ultralytics documentation and example for the same

path: /home/cvpr/CVPR/Tennis/YoloRacket/datasets/TennisRacketData/ # dataset root dir
train: images/train # train images (relative to 'path')
val: images/val # val images (relative to 'path')
test: # test images (optional)

#Keypoints
kpt_shape: [6,3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
flip_idx: [0,1,2,4,3,5]

#Classes dictionary
names:
0: tennis_racket

[glenn-jocher]

@Nit-Rathore hi there! 👋 It looks like you're encountering an issue with running inference using a TensorRT engine generated from a custom-trained YOLOv8 model. From the error screenshot and the details you've provided, it seems there might be a discrepancy between the model's expected input size or configuration and the actual input provided during inference.

Given that your custom model was trained with an input size of 640 and on a specific dataset, ensure that the same input size is used when exporting the model to TensorRT and during inference. Here's a quick reminder on how to export and run inference with the specified input size:

from ultralytics import YOLO

# Load your custom trained model
model = YOLO('path/to/your/custom_model.pt')

# Export the model to TensorRT format specifying the input size
model.export(format='engine', imgsz=640)

# Load the exported model for inference
model = YOLO('path/to/your/custom_model.engine')

# Run inference
results = model.predict('path/to/your/video.mp4')

Also, ensure that the dataset's specific configurations, like kpt_shape and flip_idx for keypoints, are correctly handled during both training and inference. If these configurations are unique to your dataset and model, they need to be consistent across all stages.

If the problem persists, it might be helpful to share the exact error message you're encountering during inference for more targeted assistance. Additionally, verifying compatibility between TensorRT 8.6.1, CUDA 12.2, and cudNN 9.0.0 with your specific model setup could also be beneficial.

Let's get this sorted! 🚀

[Nit-Rathore]

I tried the same and looked into everything, their seems to be an issue and this is the exact error message that I have been getting again and again.

Code:

import time
from ultralytics import YOLO

model = YOLO('models/tennis_racket_detect.pt')
model.export(format='engine',imgsz=640)

tensorrt_engine_path = "models/tennis_racket_detect.engine"
video_path = "test_videos/3.mp4"

--------------------------------------Exact error message--------------------------------------------

Traceback (most recent call last): File "/home/cvpr/Inference_time.py", line 12, in <module> results = tensorrt_model.predict(video_path) File "/home/cvpr/miniconda3/envs/tensorRT/lib/python3.10/site-packages/ultralytics/engine/model.py", line 439, in predict return self.predictor.predict_cli(source=source) if is_cli else self.predictor(source=source, stream=stream) File "/home/cvpr/miniconda3/envs/tensorRT/lib/python3.10/site-packages/ultralytics/engine/predictor.py", line 206, in __call__ return list(self.stream_inference(source, model, *args, **kwargs)) # merge list of Result into one File "/home/cvpr/miniconda3/envs/tensorRT/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 35, in generator_context response = gen.send(None) File "/home/cvpr/miniconda3/envs/tensorRT/lib/python3.10/site-packages/ultralytics/engine/predictor.py", line 308, in stream_inference s += self.write_results(i, self.results, (p, im, im0)) File "/home/cvpr/miniconda3/envs/tensorRT/lib/python3.10/site-packages/ultralytics/engine/predictor.py", line 173, in write_results log_string += result.verbose() File "/home/cvpr/miniconda3/envs/tensorRT/lib/python3.10/site-packages/ultralytics/engine/results.py", line 327, in verbose log_string += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " KeyError: 1

[glenn-jocher]

Hey there! 👋 It seems like the issue you're encountering is related to a KeyError during inference. This error usually occurs when there's a mismatch or an undefined class index in your model's class names.

From the snippet you've shared, the KeyError: 1 suggests that the model is predicting an object with a class index of 1, but it can't find this class index in the names dictionary provided or inferred from the model.

A potential fix could be ensuring that the names attribute of your model correctly maps all class indices to their respective class names. If you trained a custom model, make sure the names list in your dataset’s YAML file includes all classes and matches the order expected by the model.

Here's a quick check you could do:

# Check model's class names are correctly set
print(model.names)

Make sure the class names output includes the class for index 1 and corresponds to your dataset labels. If the issue persists, consider verifying the dataset's YAML file used during the model's training phase to ensure it matches the expected class indices and names.

I hope that helps! Let us know how it goes. 🚀

[Nit-Rathore]

I think this was the error, it wasn't installation as such but the training was the issue

The output displayed for model.names is for index 0 only which is tennis_racket
{0: 'tennis_racket'}

But the YAML FILE below is also having class '0' as racket:-

path: /home/cvpr/CVPR/Tennis/YoloRacket/datasets/TennisRacketData/ # dataset root dir
train: images/train # train images (relative to 'path') 4 images
val: images/val # val images (relative to 'path') 4 images
test: # test images (optional)

#Keypoints
kpt_shape: [6,3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
flip_idx: [0,1,2,4,3,5]

#Classes dictionary
names:
0: tennis_racket

[glenn-jocher]

Hello! Thanks for sharing these details. It seems like the issue is related to the class label configuration in your YAML file. The model.names output showing only {0: 'tennis_racket'} indicates that the model is recognizing only one class, which is consistent with your YAML file configuration.

Given your configuration and the model.names output, if your dataset only contains the 'tennis_racket' class and that's the intended behavior, then it appears your setup is correct for a single-class detection task. However, if you’re expecting more classes and they’re not showing up, please ensure that your dataset annotation files properly reference the class IDs matching your YAML file 'names' configuration.

For a single class scenario like yours, your setup looks fine. But remember, if expanding to multiple classes, the names section in your YAML file should include all classes, each with a unique index, like so:

names:
  0: tennis_racket
  1: another_class
  # add more classes as needed

And your annotation files should properly reference these class indexes.

If you're still encountering issues, could you provide some examples of your annotation file content? This can help identify if there's any mismatch in class labels between the dataset and the YAML configuration. Also, double-checking the dataset paths and structure could be helpful to ensure the model is training on the correct data.

Keep up the great work! 🚀

[Yonglin5170]

Hi, I encountered the same issue. The solution I found is to specify task='segment' when defining the YOLO model. I suspect the cause of this error is that different tasks have distinct parsing methods for model results. Specifying the wrong task may lead to incorrect parsing of class indices.

[Nit-Rathore]

yeah that worked, thank you!

[glenn-jocher]

@Nit-Rathore you're welcome! If you have any more questions or need further assistance, feel free to ask.