Replacing C3k2 with RepViT in YOLO11 While Reusing Pretrained Weights

Support
, , ,
1

I am exploring waste detection and I want to replace the C3k2 blocks in the YOLO11 backbone with RepViT modules, while keeping the original YOLO11 neck and detection head.

My main goal is to reuse the pretrained weights from both RepViT and YOLO11. Specifically, I would like to load the pretrained ImageNet weights for the RepViT modules and reuse the pretrained YOLO11 weights for the compatible parts of the neck and detection head, rather than training the entire model from scratch.

What is the recommended way to implement this modification in the current Ultralytics framework?

Specifically:

  1. Where should I define the RepViT module?

  2. How should I register it so it can be used in the YAML?

  3. How should I handle the different feature-map channels and strides between the RepViT modules and the original C3k2 blocks?

  4. Can I load the pretrained ImageNet weights into the RepViT modules while simultaneously reusing the compatible pretrained YOLO11 weights for the neck and detection head?

  5. If replacing C3k2 with RepViT causes channel or tensor-shape mismatches in the YOLO11 neck, what is the recommended way to handle them?

  6. Is there a recommended YAML structure for replacing the C3k2 blocks with RepViT while preserving the rest of the YOLO11 architecture?

2

The model YAML guide answers most of these questions:

How should I handle the different feature-map channels and strides between the RepViT modules and the original C3k2 blocks?

Feature map channels don’t matter. YAML builds the correct input channel shape for neck based on output channel size of the input layer. You just need to have correct c2 parsing logic in parse_model. Strides need to be the same as original.

Can I load the pretrained ImageNet weights into the RepViT modules while simultaneously reusing the compatible pretrained YOLO11 weights for the neck and detection head?

It doesn’t make sense to load the neck weights from YOLO11 because the neck weights are downstream weights that’s trained to work with backbone. It’s like loading random weights. The weights don’t have any meaning if you change the backbone.

If replacing C3k2 with RepViT causes channel or tensor-shape mismatches in the YOLO11 neck, what is the recommended way to handle them?

I don’t think there would be any general guide for how to fix shape mismatch. It’s completely dependent on the module and the mismatch.