Visualization of feature maps

[xinxin342]

🚀 Feature

Maybe we can add one more function to Yolov5 that can conveniently show the feature maps of convolutional layers and detection layers?

Motivation

Pitch

So we can visually see where the model's attention is.

Alternatives

Additional context

[github-actions]

Hello @xinxin342, thank you for your interest in our work! Please visit our Custom Training Tutorial to get started, and see our Jupyter Notebook , Docker Image, and Google Cloud Quickstart Guide for example environments.

If this is a bug report, please provide screenshots and minimum viable code to reproduce your issue, otherwise we can not help you.

If this is a custom model or data training question, please note Ultralytics does not provide free personal support. As a leader in vision ML and AI, we do offer professional consulting, from simple expert advice up to delivery of fully customized, end-to-end production solutions for our clients, such as:

• Cloud-based AI systems operating on hundreds of HD video streams in realtime.
• Edge AI integrated into custom iOS and Android apps for realtime 30 FPS video inference.
• Custom data training, hyperparameter evolution, and model exportation to any destination.

For more information please visit https://www.ultralytics.com.

[glenn-jocher]

@xinxin342 we are open to PRs, so if you would like to work on this and submit a PR we'd be happy to review!

[xinxin342]

@ xinxin342我们对PR保持开放，因此，如果您要从事此工作并提交PR，我们将很高兴进行审查！

Sorry, I'm completely out of my depth, but I see that it looks like YOLOv4 can do this.

[glenn-jocher]

@xinxin342 feature map visualization is something you can inject into the model forward function at the location of your choosing. It's up to you exactly where in the model you'd want to visualize this, as there are hundreds of layers and thousands of maps.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[glenn-jocher]

@xinxin342 good news 😃! Feature map visualization was added ✅ in PR #3804 by @Zigars today. This allows for visualizing feature maps from any part of the model from any function (i.e. detect.py, train.py, test.py). Feature maps are saved as *.png files in runs/features/exp directory. To turn on feature visualization set feature_vis=True in the model forward method and define the layer you want to visualize (default is SPP layer).

yolov5/models/yolo.py

Lines 158 to 160 in 20d45aa

	if feature_vis and m.type == 'models.common.SPP':
	feature_visualization(x, m.type, m.i)

To receive this update:

• Git – git pull from within your yolov5/ directory or git clone https://github.com/ultralytics/yolov5 again
• PyTorch Hub – Force-reload with model = torch.hub.load('ultralytics/yolov5', 'yolov5s', force_reload=True)
• Notebooks – View updated notebooks
• Docker – sudo docker pull ultralytics/yolov5:latest to update your image

Thank you for spotting this issue and informing us of the problem. Please let us know if this update resolves the issue for you, and feel free to inform us of any other issues you discover or feature requests that come to mind. Happy trainings with YOLOv5 🚀!

[MLDavies]

Can you tell me more about how to implement the feature maps functionality? I assumed from above that, if we're after the SPP layer, we just needed to go into yolo.py and change visualize=False to True in: def _forward_once(self, x, profile=False, visualize=True). Clearly this isn't enough. What am I missing?

[glenn-jocher]

@MLDavies 👋 Hello! Thanks for asking about feature visualization. YOLOv5 🚀 features can be visualized through all stages of the model from input to output. To visualize features from a given source run detect.py with the --visualize flag:

python detect.py --weights yolov5s.pt --source data/images/bus.jpg --visualize

An example Notebook visualizing bus.jpg features with YOLOv5s is shown below:

All stages are visualized by default, each with its own PNG showing the first 32 feature maps output from that stage. You can open any PNG for a closer look. For example the first 32 feature maps of the Focus() layer output are shown in stage0_Focus_features.png:

Feature maps may be customized by updating the feature_visualization() function in utils/plots.py:

yolov5/utils/plots.py

Lines 403 to 427 in bb5ebc2

	def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detect/exp')):
	"""
	x: Features to be visualized
	module_type: Module type
	stage: Module stage within model
	n: Maximum number of feature maps to plot
	save_dir: Directory to save results
	"""
	if 'Detect' not in module_type:
	batch, channels, height, width = x.shape # batch, channels, height, width
	if height > 1 and width > 1:
	f = f"stage{stage}_{module_type.split('.')[-1]}_features.png" # filename
	blocks = torch.chunk(x[0].cpu(), channels, dim=0) # select batch index 0, block by channels
	n = min(n, channels) # number of plots
	fig, ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True) # 8 rows x n/8 cols
	ax = ax.ravel()
	plt.subplots_adjust(wspace=0.05, hspace=0.05)
	for i in range(n):
	ax[i].imshow(blocks[i].squeeze()) # cmap='gray'
	ax[i].axis('off')
	print(f'Saving {save_dir / f}... ({n}/{channels})')
	plt.savefig(save_dir / f, dpi=300, bbox_inches='tight')
	plt.close()

Good luck 🍀 and let us know if you have any other questions!

[MLDavies]

Perfect. Love yolov5! Ya, I was (mistakenly) trying to call --visualize on train.py rather than detect.py. Thanks for what you do, and for your responses.

[wanglegoc]

@glenn-jocher I run the following script python detect.py --weights yolov5s.pt --source data/images/bus.jpg --visualize. Instead of obtaining results as you shown, here how it looks for stage22_Concat_features

[glenn-jocher]

@wanglegoc --visualize works correctly. I just tested right now:

[abhigoku10]

@glenn-jocher can we use this feature for yolov4 architecture ? if not what are the changes to be made to the code

[glenn-jocher]

@abhigoku10 The feature map visualization is specific to the YOLOv5 architecture and is not directly compatible with YOLOv4. Porting this feature to YOLOv4 would require significant modifications to the codebase, including changes to the network architecture, model forward method, and visualization logic. If you are interested in adding feature map visualization to YOLOv4, I recommend referring to the YOLOv5 implementation as a reference and making the necessary adaptations to the YOLOv4 codebase. Good luck!