PyTorch特征图的可视化

说明:

  1. max_num表示将存储的某个特征图通道数,因此max_num不能大于输入的特征图的通道维度
  2. x表示输入的特征图
  3. out_path表示输出的文件夹
  4. self因为本人使用的时候是放在backbone中的所以这个删掉即可
import os
import cv2
import numpy as np

def vis_feature(self, x, max_num=5, out_path='/home/XX/'):
    for i in range(0, x.shape[1]):
        if i >= max_num:
            break
        feature = x[0, i, :, :].view(x.shape[-2], x.shape[-1])
        feature = feature.cpu().numpy()
        feature = 1.0 / (1 + np.exp(-1 * feature))
        feature = np.round(feature * 255)
        feature_img = cv2.applyColorMap(feature, cv2.COLORMAP_JET)
        dst_path = os.path.join(out_path, i + '.png')
        cv2.imwrite(dst_path, feature_img)

你可能感兴趣的:(pytorch,pytorch,深度学习)