神经网络迭代时可视化中间特征

2021-11-17

字数统计: 517字 | 阅读时长: 2分

阅读量

神经网络迭代时，在forward中可视化中间特征，用到Image库。（其实pytorch的transform库也可以）

将形状为(b,c,h,w)的Tensor可视化为图像的具体步骤：

选取每个batch的第一张图像进行可视化
考虑经过卷积后，中间的特征层维度较高，只可视化前10维
对每一维：均值归一化，再变换回自然图像，Image库保存到本地

# feature_map：Tensor(b,c,h,w)
def visualize(feature_map):
    feature_map = feature_map[0]
    feature_map = feature_map.asnumpy()
    num_channel = feature_map.shape[0]
    if num_channel>=10: # 最多展示10层
        num_channel = 10
    for i in range(num_channel):
        feature_image = feature_map[i, :, :]
        feature_image-= feature_image.mean()
        feature_image/= feature_image.std ()
        feature_image*=  64
        feature_image+= 128
        pil_img = Image.fromarray(feature_image)
        pil_img = pil_img.convert('RGB')
        pil_img.save('visualize_imgs{}.png'.format(i))
        
    return 0

在模型的forward中加入此函数即可。

(注：在Mindspore中试验过，不确定pytorch中取feature_map[0]会不会自动降一维，待补充。)

实际使用中，我加入时间戳标记这一组图片。

def forward(self, x):
    time_tuple = time.localtime(time.time())
	# ......
    visualize(x[0],0,time_tuple)
    x = self.conv1(x)
    visualize(x[0],1,time_tuple)
    # ......

def visualize(feature_map, loc, time_tuple):
    feature_map = feature_map.asnumpy()
    num_channel = feature_map.shape[0]
    if num_channel>=10: # 最多展示10层
        num_channel = 10
    for i in range(num_channel):
        feature_image = feature_map[i, :, :]
        feature_image-= feature_image.mean()
        feature_image/= feature_image.std ()
        feature_image*=  64
        feature_image+= 128
        pil_img = Image.fromarray(feature_image)
        pil_img = pil_img.convert('RGB')
        pil_img.save('visualize_imgs/{}{}{}_{}_{}.jpg'.format(time_tuple[3],time_tuple[4],time_tuple[5],loc,i))
        
    return 0

绘图所用函数，用到matplotlib：

import matplotlib.pyplot as plt
import glob
from PIL import Image
import numpy as np

imgs_path = glob.glob('visualize_imgs/*')
sort_imgs_dict = {}
for img_path in imgs_path:
    mark = img_path.split('\\')[1]
    mark = mark.split('_')[0]
    # print(mark)
    if mark in sort_imgs_dict:
        sort_imgs_dict[mark].append(img_path)
    else:
        sort_imgs_dict[mark] = [img_path]
# print(sort_imgs_dict)
plt.subplots_adjust(hspace=0.2, wspace=0.1)
for key in sort_imgs_dict.keys():
    each_imgs_path = sort_imgs_dict[key]
    print('num:', len(each_imgs_path), 'key', key)
    n = 0
    for i in range(3):
        n = n + 1
        img = np.array(Image.open(each_imgs_path[n-1]))
        plt.subplot(9, 10, n)
        plt.savefig('vis_{}.png'.format(key))
        plt.imshow(img)
        plt.axis('off')

    for i in range(80):
        n = n + 1
        img = np.array(Image.open(each_imgs_path[n-1]))
        plt.subplot(9, 10, n+7)
        plt.savefig('vis_{}.png'.format(key))
        plt.imshow(img)
        plt.axis('off')

效果：