Tensorboard 图片特征图权重参数可视化

Posted 2022-09-27 洪流之源

1. 可视化单张图片

from torch.utils.tensorboard import SummaryWriter
import torch

if __name__ == '__main__':

    summary_writer = SummaryWriter(log_dir='log_image', comment='test tensorboard image', filename_suffix='test_1')

    image1 = torch.randn(3, 640, 640)

    # 可视化图片
    # tag : 为图像的名称
    # img_tensor : 可视化的tensor数据(图像数据), 为了方便可视化, 会对图像数据像素值有选择的缩放到0~255
    #              如果图像数据像素值都不大于1，则可视化时所有的像素值会乘以255
    #              如果存在像素值大于1的像素值，则可视化时不会对所有像素乘以255
    # global_step : 可视化图像横轴的索引，如果多个图片可进行鼠标拖动显示不同的图片
    # walltime : 用于记录发生的时间，默认为 time.time()
    # dataformats : 图像数据排布, CHW、HWC、HW, 默认为CHW
    # image1的值为正态分布数据，图像数据像素值存在大于1的值，则可视化时所有的像素值不会会乘以255，所以可视化都是原始的正态分布数据的画面
    summary_writer.add_image(tag='image1', img_tensor=image1, global_step=1, walltime=None, dataformats="CHW")

    image2 = torch.ones(3, 640, 640)
    # image2的值都为1，图像数据像素值都不大于1，则可视化时所有的像素值会乘以255，所以可视化都是白色的画面
    summary_writer.add_image(tag='image2', img_tensor=image2, global_step=2, walltime=None, dataformats="CHW")

    summary_writer.close()

运行tensorboard，并指定日志目录：

tensorboard --logdir log_image

2.  通过make_grid显示多张图片

from torch.utils.tensorboard import SummaryWriter
import torch
import torchvision

if __name__ == '__main__':

    summary_writer = SummaryWriter(log_dir='log_grid', comment='test tensorboard image', filename_suffix='test_1')

    images = torch.randn(4, 3, 640, 640)

    # 制作图像网格
    # tensor : 图像数据(tensor数据)，N * C * H * W形式
    # nrow : 行数 (列数自动计算), 比如16张图像, nrow设置为4, 则列数为4
    # padding : 图像间距(像素单位)
    # normalize : 是否将像素标准化, 是否将数据缩放到0 ~ 255
    # value_range : 标准化范围, 如果原始数据的范围在[-100， 200], range=[-50, 200], 
    #               则会对原始数据按照range的范围进行截断(小于-50的都设置为-50)，然后再进行标准化(缩放到0 ~ 255)
    # scale_each : 是否单张图维度标准化
    # pad_value : 图像间距的像素值
    images_grid = torchvision.utils.make_grid(tensor=images, nrow=2, padding=10, 
                                              normalize=False, value_range=None, 
                                              scale_each=False, pad_value=0)

    
    summary_writer.add_image(tag='image', img_tensor=images_grid, global_step=0)

    summary_writer.close()

运行tensorboard，并指定日志目录：

tensorboard --logdir log_grid

如下，4张图片显示： 

 3. 可视化特征图

读取lena.jpg图片，输入到alexnet第一个卷积层，对得到的特征图可视化：

from torch.utils.tensorboard import SummaryWriter
import torchvision
from torchvision import transforms

if __name__ == '__main__':
    summary_writer = SummaryWriter(log_dir='feature_map')

    norm_transform = transforms.Normalize([0.49139968, 0.48215827, 0.44653124], 
                                          [0.24703233, 0.24348505, 0.26158768])

    image_transforms = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        norm_transform
    ])

    from PIL import Image
    image = Image.open("./lena.jpg").convert('RGB')
    if image_transforms is not None:
        image_tensor = image_transforms(image)
    image_tensor.unsqueeze(0)    # chw --> bchw

    # 模型
    alexnet = torchvision.models.alexnet(pretrained=True)
    print(alexnet)

    # forward
    conv_layer_1 = alexnet.features[0]
    feature_map_1 = conv_layer_1(image_tensor)
    print(feature_map_1.shape)

    # 预处理
    feature_map_1 = feature_map_1.unsqueeze(0)  # bchw=(1, 64, 55, 55) --> (64, 1, 55, 55)
    feature_map_1 = feature_map_1.transpose(0, 1)
    
    fmap_1_grid = torchvision.utils.make_grid(feature_map_1, normalize=True, scale_each=True, nrow=8)
    summary_writer.add_image('feature map in conv1', fmap_1_grid, global_step=322)
    summary_writer.close()

运行tensorboard，并指定日志目录：

tensorboard --logdir feature_map

 4. 可视化权重参数

可视化alexnet第一个卷积层的权重参数：

from torch.utils.tensorboard import SummaryWriter
import torch.nn as nn
import torchvision

if __name__ == '__main__':

    summary_writer = SummaryWriter(log_dir='weight')

    alexnet = torchvision.models.alexnet(pretrained=True)

    for index, module in enumerate(alexnet.modules()):
        if isinstance(module, nn.Conv2d):
            kernels = module.weight
            kernel_nums, kernel_channels, kernel_width, kernel_height = tuple(kernels.shape)
            print(kernel_nums, kernel_channels, kernel_width, kernel_height)

            # 单独可视化每个kernel
            for kernel_id in range(kernel_nums):
                kernel = kernels[kernel_id, :, :, :].unsqueeze(1)   # make_grid需要 BCHW，这里拓展C维度
                kernel_grid = torchvision.utils.make_grid(kernel, normalize=True, scale_each=True, nrow=kernel_channels)
                summary_writer.add_image('_Convlayer_split_in_channel'.format(index), kernel_grid, global_step=kernel_id)

            # 可视化当前层所有的kernel
            kernel_all = kernels.view(-1, 3, kernel_height, kernel_width)  # 3, h, w
            kernel_grid = torchvision.utils.make_grid(kernel_all, normalize=True, scale_each=True, nrow=8)  # c, h, w
            summary_writer.add_image('_all'.format(index), kernel_grid, global_step=322)

            break

    summary_writer.close()

运行tensorboard，并指定日志目录：

tensorboard --logdir weight