yolov5特征图可视化

看到一篇帖子讲可视化,使用的效果很棒,特意记下来以防迷路。会给出原文链接,看到的UU们可以去看原文更详细。

使用方法

在utils中的general.py或者plots.py添加如下函数:(我添加在general中,添加在其他位置注意加引用)

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)

接着在models中的yolo.py中添加代码(add in here):

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
            feature_vis = True
            if m.type == 'models.common.C3' and feature_vis:
                print(m.type, m.i)
                feature_visualization(x, m.type, m.i)

 
        if profile:
            print('%.1fms total' % sum(dt))
        return x

添加在yolo.py后,无论是在detect.py还是在train.py中都会进行可视化特征图。

然而训练的过程中并不一定需要一直可视化特征图,feature_vis参数是用来控制是否保存可视化特征图的,保存的特征图会存在features文件夹中。如果想看其它层的特征只需要修改m.type或是用m.i来进行判断是否可视化特征图。m.type对应的是yaml文件中的module,即yolov5的基础模块,例如c3,conv,spp等等,而m.i则更好理解,即是模块的id,通常就是顺序,如果你尝试修改过配置文件,那么你肯定知道是什么。
————————————————
版权声明:本文为CSDN博主「Silence_Zzz」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/qq_40231159/article/details/118270178

你可能感兴趣的:(yolov5,python,机器学习,深度学习)