[RFC] Support YOLOX detection model

[zhiqwang]

🚀 The feature

YOLO aka. You Only Look Once, which is a vibrant series of object detection models since the release of Joseph Redmon You Only Look Once: Unified, Real-Time Object Detection.

So far a couple of more notable implementations are as follows (all PyTorch):

• YOLOv3 2018. Cited by 14369 ¹²
• YOLOv4 2020. Cited by 4905
• YOLOv5 2020. Starred at GitHub 29.4k ³
• YOLOX 2021. Cited by 299 ⁴
• YOLOv7 2022. ⁵

Motivation, pitch

Until now, one of the most successful ones is probably YOLOv5. YOLOv5 is great, and they have also built up a very friendly community and ecosystem. We don't intend to copy YOLOv5 into TorchVision, our main goal here is to make training SoTA models easier and share reusable subcomponents to build the next SoTA models in the same/proxy family.⁶

YOLOX is a high-performance anchor-free YOLO, and it has a good balance in terms of copyright and code quality, it's enough to have a YOLOX implementation from the community's perspective.

The License

YOLO{v5/v7} are built under the GPL-3.0 license, and YOLOX is built under the Apache-2.0 license.

More context

I have previously rewritten the code used in the inference part of YOLOv5 according to the style and specification of torchvision⁷, and I can relicense that part to BSD-3-Clause license. The amount of work involved in the model inference part is not much with the help of YOLOX base code.

Data augmentation and a new trainer engine will be the core of what we will do here.

The data augmentation section is in the planning list #6224 , and we have already merged some augmentation methods like #5825 , I think it would help us to build the next SoTA models with a new primitives like classification models.⁸

• Mosaic
• augment_hsv
• random_perspective
• mixup
• random_flip, left-right

As TorchVision adds more and more models, it may be time to abstract out a simple trainer engine for sharing reusable subcomponents. It might be more appropriate to open a new thread for necessity and specific steps about this part.

cc @datumbox @YosuaMichael @oke-aditya

Footnotes

1. https://github.com/ultralytics/yolov3/tree/v9.1 ↩

2. https://github.com/open-mmlab/mmdetection/tree/master/configs/yolo ↩

3. https://github.com/ultralytics/yolov5 ↩

4. https://github.com/Megvii-BaseDetection/YOLOX ↩

5. https://github.com/WongKinYiu/yolov7 ↩

6. https://github.com/keras-team/keras-cv/issues/622#issuecomment-1198063712 ↩

7. https://github.com/zhiqwang/yolov5-rt-stack/tree/main/yolort/models ↩

8. https://pytorch.org/blog/how-to-train-state-of-the-art-models-using-torchvision-latest-primitives/ ↩

[datumbox]

@zhiqwang Thank you very much for the comprehensive proposal. :)

Your implementation at YOLO5-RT-stack is indeed of very high quality. Having a modern implementation of YOLO was on our bucket list but I just want to be mindful and not cannibalise your project. After all PyTorch's unique value proposition is its rich ecosystem. Having said that, if you are happy to upstream parts of your repo to TorchVision then we would absolutely love to have it. Rest assured that if we do add it, we will make sure to provide all the necessary credit to OSS contributors who made that possible. I know that your coding-styles and practices are very aligned with the ones used in TorchVision, so I agree v5 would probably be the easiest step forward.

I have a couple of questions for you:

1. Did you manage to reproduce the original results using your own implementation or you provide ported weights? Do you have access to a cluster to train such models (don't worry if you don't, we can work something out on our side)?
2. Other than the transforms you mentioned above which we already added/plan-to-add, what other augmentations are we missing? I suspect Mosaic was used for this model. Anything else?
3. I assume you will implement CSPDarknet for the backbone. Which other backbones do you plan to drop on the classification side?
4. Can you clarify whether your implementation is a complete rewrite or if you use components from a GPL3 implementation? If you didn't fully rewrite the code, it would be safer for us to reach out to the authors of the reference implementations and confirm they would be happy for us to progress.

Concerning the training engine, I complete agree we should refactor large part of our reference scripts to inherit and reuse components. My recommendation though is not to link this work to the addition of YOLO as this is already a very big project. There are also various potential solutions that we might want to leverage (for example TorchRecipes), but this will require additional chats.

I would suggest next steps is to clarify the above and decide how to progress. Possibly we will need to split the project to subprojects and potentially invite more contributors to help out. We could tackle this as part of #6323 and leverage the community.

[zhiqwang]

Hi @datumbox

1. Did you manage to reproduce the original results using your own implementation or you provide ported weights? Do you have access to a cluster to train such models (don't worry if you don't, we can work something out on our side)?

I have only verified on a few images with the ported weights, because there are some differences in our preprocessing part compared to the original yolov5 version, I have not been able to do a complete verification on the coco dataset, and I can add some more detailed comparison data in the next few days.

I don't have a server in my hand. We can divide the task into smaller parts and eventually put all the modules together, and then being able to train with your help is the best option.

2. Other than the transforms you mentioned above which we already added/plan-to-add, what other augmentations are we missing? I suspect Mosaic was used for this model. Anything else?

As you said the most important part of the data augmentation is the mosaic technique. It was first introduced in ultralytics/yolov3#310 (comment), and there is a similar discussion in AlexeyAB/darknet#3114 (comment). This mosaic technique is helpful to detect smaller size object. (And I think this is the key technology that allows the YOLO series to be trained from scratch.) I quote Jocher's conclusions below.

The smaller cars are detected earlier with less blinking and cars of all sizes show better behaved bounding boxes.

In addition, v5 also uses the following enhancements:

• augment_hsv
• random_perspective
• mixup
• random_flip, left-right

Seems that the random_fliplr augmentation is the RandomHorizontalFlip in TorchVision:

vision/references/detection/transforms.py

Line 30 in b30fa5c

class RandomHorizontalFlip(T.RandomHorizontalFlip):

[zhiqwang]

3. I assume you will implement CSPDarknet for the backbone. Which other backbones do you plan to drop on the classification side?

YOLO{v5/v7/X} train their detection models from scratch, and they now have backbone that is not the DarkNet presented in the original paper. It would be nice to have a pre-trained model on the ImageNet to help accelerate our training, but it is up for debate whether it is necessary to implement the original version of CSPDarknet.

4. Can you clarify whether your implementation is a complete rewrite or if you use components from a GPL3 implementation? If you didn't fully rewrite the code, it would be safer for us to reach out to the authors of the reference implementations and confirm they would be happy for us to progress.

The codes currently in this folder https://github.com/zhiqwang/yolov5-rt-stack/tree/main/yolort/models was written from scratch, I only called some common functions of YOLOv5. That parts have been rewritten by YOLOX, we can call YOLOX's common functions instead (or we can rewrite it ourselves) to get rid of this dependency. The main reason I used YOLOv5's common was to be able to load the checkpoints trained from YOLOv5.

Concretely I restructured the YOLOv5's yaml-parse mechanism into following three sub-modules in the layout of TorchVision:

• BackboneWithPAN
• YOLOHead
• PostProcess

[datumbox]

@zhiqwang I had the chance to investigate a bit further the references. The biggest concern about YOLOv5 is that there is still no paper to accompany the architecture (see ultralytics/yolov5#1333); I remember that it first came out as a repo and the owners said that the paper will be coming out shortly but I don't think there is currently one. Though it's a very popular architecture which achieves good results, the lack of paper is a problem as we usually focus on canonical implementations and expansions that have been studied in research. YOLOX seems like a viable alternative. Perhaps that's the way forward to avoid licensing issues, wdyt?

I have not been able to do a complete verification on the coco dataset, and I can add some more detailed comparison data in the next few days.

Sounds good, I think it's worth confirming that the implementation yields the expected accuracy prior deciding to adopt it.

then being able to train with your help is the best option

Sounds good. We can follow a similar approach as with FCOS but with omitting the original training. We have capacity of training such network internally, so you don't have to have your own infra.

mosaic

We should implement the mosaic augmentation and add it on references at first. Then once @pmeier and @vfdev-5 are back, we can examine implementing them as transforms on the new API.

mixup

I'm a bit surprised to see it in the list (haven't check the references). Do they use mixup for detection? Do they adjust the probabilities of the labels? How about the boxes, are they multiplied also by weights?

It would be nice to have a pre-trained model on the ImageNet to help accelerate our training, but it is up for debate whether it is necessary to implement the original version of CSPDarknet.

I can confirm that training from scratch usually yields better results. I'm OK not adding a Darknet arch in TorchVision; it's kind of old.

[oke-aditya]

I can confirm that training from scratch usually yields better results. I'm OK not adding a Darknet arch in TorchVision; it's kind of old.

Agreed but as such we will need DarkNet to implement YOLOX. So it might just be better to provide it too. Although it's old it's still being used and relevant. It's quite a fundamental model like AlexNet. So somehow I feel it might be good to add.

[zhiqwang]

YOLOX seems like a viable alternative. Perhaps that's the way forward to avoid licensing issues

Hi @datumbox , agree on this one, It is a good choice to add YOLOX to the list first.

[datumbox]

@zhiqwang Sounds good. I'll need to dig a bit on the YOLOX paper and familiarize myself. I'll try to do this by EOW. In meantime, if you have in your mind a clear plan on intermediary milestones for adding YOLOX please add it here (aka addition of X arch, Y transforms, Z operators etc). This will hopefully let us to coordinate among contributors.

[datumbox]

@zhiqwang I'm late by 1 week. Sorry I got caught up on other pieces of work.

I've gone through the bibliography around YOLOX and here are some thoughts:

• As noted earlier we will need to implement Mosaic Transform.
• Extend Mixup for Object Detection (ref - section 3.1) cc @pmeier @vfdev-5
• Implement DarkNet53 and/or Modified CSPNet. This depends on which variant sizes we want to offer. IMO, especially with modified CSPNet, I would probably want to avoid offering it as a standalone classification model of Vision.
• YOLOX is more based on YOLOv3 rather than YOLOv5. This has the positive that it's based on published work but negative that it is less relevant to previous work you did on the YOLO5-RT-stack repo.

Mosaic and MixUp are worth implementing. I've added tasks for them on the #6323 issue. Whether we will go ahead with the rest depends on your bandwidth. Is this something you would like to pick up and lead? If yes, we can find a POC on our side that would assist with the model validation and training resources. Let me know, thanks!

[zhiqwang]

Hi @datumbox ,

YOLOX is more based on YOLOv3 rather than YOLOv5. This has the positive that it's based on published work but negative that it is less relevant to previous work you did on the YOLO5-RT-stack repo.

I agree with you here. YOLOX have a good balance in terms of copyright and code quality, and it's enough to have a YOLOX implementation from the community's perspective.

Whether we will go ahead with the rest depends on your bandwidth. Is this something you would like to pick up and lead?

Sorry for not having enough bandwidth to work here recently ( But I can help to review the codes and support for deployment if there is such a need.

[datumbox]

Sorry for not having enough bandwidth to work here recently ( But I can help to review the codes and support for deployment if there is also a need for this.

@zhiqwang Thanks for getting back to me. I completely understand. Unfortunately we are very constrained in terms of headcount and bandwidth at the moment. I don't think any of the maintainers can pick this up. Originally the idea of you picking up and leading this initiative was very promising as you have extensive experience with the YOLO architecture due to your earlier work. But I understand that since we will be interested in porting YOLOX and not v5, that would increase significantly your work. I'm happy to leave this open in case your situation changes on the future.

Since we are here, let me do the cheeky move and check if any of the original authors of YOLOX would be interested in contributing an implementation to TorchVision? @FateScript @Joker316701882 @GOATmessi7