Include NMS in model export?

[mattpopovich]

Question:

I currently have a custom yolov5 model running in my C++ pipeline with TorchScript. The TorchScript model was obtained by running export.py. Getting the TorchScript model to run on the GPU in C++ is easy enough via model_gpu = torch::jit::load(model_path, torch::kCUDA);. I would now like to run all of pre-processing, inference, and post-processing on the GPU to speed up my C++ pipeline. Is there any way to export pre-processing and post-processing (non-maximum suppression aka NMS) IN the model? Or would you recommend something else for running the whole pipeline in the GPU in C++?

What I've tried:

Running the TorchScript model

When you run the TorchScript model (with size = 640 and after pre-processing), it outputs:

• Tuple[Tensor, List[Tensor]]
  • Tuple[0] = [1, 25200, 13] Tensor
  • Tuple[1] = List[Tensor] of size 3
    • Tuple[1][0] = [1, 3, 80, 80, 13] Tensor
    • Tuple[1][1] = [1, 3, 40, 40, 13] Tensor
    • Tuple[1][2] = [1, 3, 20, 20, 13] Tensor

These are all of the detections. After being put through NMS (post-processing), we get our actual (non-duplicate) detections.

Using jit.script instead of jit.trace

In export.py, I can add:

sm = torch.jit.script(model(img))    # ScriptModule
sm.save(f)

Which gives me an exception: "Could not get name of python class object"

This isn't a big deal because jit.trace works, but it's worth noting this error here, because I get it below as well.
This might be because TorchScript (jit.script) only supports PyTorch and the math module. There are a few numpy calls in AutoShape.forward() so maybe that is one of the issues preventing this?

Exporting the PyTorch Hub model

It appears that the PyTorch Hub model includes both pre-processing and post-processing (NMS) via the AutoShape class. My thought was to import my custom model using torch.hub.load(model='custom', [...]) (which is successful), then I would either jit.trace or jit.script in order to take advantage of the built-in pre and post-processing.

• jit.trace of torch.hub model
  • Comparison exception: With rtol=1e-05 and atol=1e-08, found 1 element(s) (out of 1) whose difference(s) exceeded the margin of error (including 0 nan comparisons). The greatest difference was 1.5 (2.0 vs. 0.5), which occurred at index 0.
• jit.script of torch.hub model
  • RuntimeError: Could not get name of python class object

So.. looks like this isn't a viable alternative either.

Appreciate any advice!

[zhiqwang]

Hi @mattpopovich , I've done some experiments about exporting the letterbox (pre-processing) and NMS (post-processing) in AutoShape to the torchscript using torch.jit.script here.

[mattpopovich]

@zhiqwang I did come across your repo! It is on my list of things to investigate in the future. Thanks for putting it together and giving me the heads up!

[zhiqwang]

Hi @mattpopovich

torch.jit.trace has more practical applications, torch.jit.script is more experimental than torch.jit.trace.

But script is more flexible, if you want to export the lettexbox and NMS to torchscript, seems that we must use torch.jit.script, and that's what I've done in my repo.

[github-actions]

👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs.

Access additional YOLOv5 🚀 resources:

• Wiki – https://github.com/ultralytics/yolov5/wiki
• Tutorials – https://docs.ultralytics.com/yolov5
• Docs – https://docs.ultralytics.com

Access additional Ultralytics ⚡ resources:

• Ultralytics HUB – https://ultralytics.com
• Vision API – https://ultralytics.com/yolov5
• About Us – https://ultralytics.com/about
• Join Our Team – https://ultralytics.com/work
• Contact Us – https://ultralytics.com/contact

Feel free to inform us of any other issues you discover or feature requests that come to mind in the future. Pull Requests (PRs) are also always welcomed!

Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐!