Remove detection head

[Neloy262]

Search before asking

• I have searched the YOLOv5 issues and discussions and found no similar questions.

Question

I have trained a yolov5n model on a custom dataset. Now i want to remove the detection head so i can get the features immediately before the detection head. How might I do this?

Additional

No response

[UltralyticsAssistant]

👋 Hello @Neloy262, thank you for reaching out and for using YOLOv5 🚀! This is an automated response, and an Ultralytics engineer will assist you soon.

For now, I recommend starting with our ⭐️ Tutorials to explore various tasks related to custom model training and configuration.

For your question about modifying the model architecture, including removing parts like the detection head, you might want to dive into our Custom Data Training guides.

If this is related to a 🐛 Bug, please ensure you provide a minimum reproducible example to help us debug.

For custom training configurations ❓, please include more details where possible, such as your model setup, dataset examples, and any modifications applied.

Requirements

Make sure your environment meets our requirements:

Python>=3.8.0 and all requirements.txt installed. You can start by running:

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt  # install

Helpful Environments

YOLOv5 can be developed in any verified environment:

• Notebooks with free GPUs:
• Cloud Providers: Check our GCP and AWS guides
• Docker Image: See our Docker Guide

Stay Updated

The green badge indicates passing CI tests for operations like training and validation.

Discover YOLOv8 🚀

You can also explore our latest model YOLOv8, optimized for a range of tasks and improved performance. Get started by installing:

pip install ultralytics

Check out our YOLOv8 Docs for more.

Happy coding! 😊

[glenn-jocher]

@Neloy262 to remove the detection head and access features before it, you can modify the model by excluding the final layers responsible for detection. You can achieve this by editing the model's architecture in the models/yolov5n.yaml file, removing or commenting out the Detect module. Then, load the modified model for inference. For detailed guidance, refer to the YOLOv5 architecture documentation.