a question in yolo checkpoint converter

[junghyun-avikus]

https://github.com/zhiqwang/yolov5-rt-stack/blob/8b578eb9a7910f1dcb28188a36c8c540d15a9430/yolort/models/_checkpoint.py#L53-L64

The above lines seem to be hard coded based on the Ultralytics yolov5.
I, however, made a small change in the model, and the model is defined by the following yaml.

# YOLOv5 🚀 by YOLOAir, GPL-3.0 license

# Parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5 v6.0 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 6, C3, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, C3, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 3, C3, [1024]],
   [-1, 1, SPPF, [1024, 5]],  # 9
  ]

# YOLOv5 v6.0 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, C3, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
   [-1, 1, CBAM, [1024]],

   [[17, 20, 24], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

Everything statys the same, but at the end, I added an extra CBAM module right before the Detect.
So I made a change to the head index, increased by 1, but I am not sure what block maps are for.
Should I make further changes?

[zhiqwang]

Modification is necessary, and it is quite messy actually. The main point I did was to print out the module (via print(module)) and check the mapping between the module in yolov5 and yolort, and then hard code it here. This requires you to be very familiar with both the yolov5 and yolort model building approaches. So I don't particularly recommend that you deploy your models for TensorRT in this way.

Our main effort is to link nvidia's efficientNMS plugin to the model for TensorRT as reported in here, and the recommended way is to link the plugin together by modifying the ONNX graph.