feature map visualization

[Zigars]

🚀 Feature

I find many people want get the feature map in model middle(include myself), but you did not give the specific method to visualization feature map. Also I find a previous issue and know how can I come true this. So I provide a feature_visualization function so that people can visualization feature map by using yolov5's code .

This is my effect picture :

Motivation

It's easy to use. just add feature_visualization function in utils/general.py or utils/plots.py：

import matplotlib.pyplot as plt
from torchvision import transforms
 
def feature_visualization(features, model_type, model_id, feature_num=64):
    """
    features: The feature map which you need to visualization
    model_type: The type of feature map
    model_id: The id of feature map
    feature_num: The amount of visualization you need
    """
    save_dir = "features/"
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    # print(features.shape)
    # block by channel dimension
    blocks = torch.chunk(features, features.shape[1], dim=1)
 
    # # size of feature
    # size = features.shape[2], features.shape[3]
 
    plt.figure()
    for i in range(feature_num):
        torch.squeeze(blocks[i])
        feature = transforms.ToPILImage()(blocks[i].squeeze())
        # print(feature)
        ax = plt.subplot(int(math.sqrt(feature_num)), int(math.sqrt(feature_num)), i+1)
        ax.set_xticks([])
        ax.set_yticks([])
 
        plt.imshow(feature)
        # gray feature
        # plt.imshow(feature, cmap='gray')
 
    # plt.show()
    plt.savefig(save_dir + '{}_{}_feature_map_{}.png'
                .format(model_type.split('.')[2], model_id, feature_num), dpi=300)

and than add this in yolo.py:

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

    def forward_once(self, x, profile=False):
        y, dt = [], []  # outputs
        for m in self.model:
            if m.f != -1:  # if not from previous layer
                x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]  # from earlier layers

            if profile:
                o = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 if thop else 0  # FLOPS
                t = time_synchronized()
                for _ in range(10):
                    _ = m(x)
                dt.append((time_synchronized() - t) * 100)
                print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))

            x = m(x)  # run
            y.append(x if m.i in self.save else None)  # save output

            # add in here

My code is not so concise, but I think feature map visualization function can help people understand what does the convolution operation do intuitively.

Also I have a little confuse about Model function, I modified in yolo.py, but when I run detect.py, I also can get the feature map.
detect.py use model = attempt_load(weights, map_location=device) to load model, but I can't find the relationship with yolo.py and detect.py, I'm not familiar with this mechanism in the PyTorch code. maybe you can solve my confuse!

Pitch

Alternatives

Additional context

[Zigars]

Now the code can read.

[glenn-jocher]

@Zigars thanks for the feature suggestion and your example code. Can you please submit a PR with your implementation? We will then review the PR and try to get this merged to help everyone else plot features also. See https://docs.ultralytics.com/help/contributing/ to get started. Thanks!

[Zigars]

@glenn-jocher ok, I will do it.

[glenn-jocher]

@Zigars great thanks!

[glenn-jocher]

@Zigars 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 🚀!

[glenn-jocher]

@Zigars good news 😃! Your original issue may now be fixed ✅ in PR #3920. This is a complete revamp of feature visualization with many fixes and improvements.

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 🚀!