【深度之眼】Pytorch框架班第五期-Tensorboard图像、卷积核和特征图可视化代码解析

import os
import torch
import time
import torchvision.models as models
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DataLoader

import sys
hello_pytorch_DIR = os.path.abspath(os.path.dirname(__file__)+os.path.sep+".."+os.path.sep+"..")
sys.path.append(hello_pytorch_DIR)
from tools.my_dataset import RMBDataset
from tools.common_tools import set_seed
from model.lenet import LeNet


set_seed(1)  # 设置随机种子


# ----------------------------------- 3 image -----------------------------------
flag = 0
# flag = 1
if flag:

    writer = SummaryWriter(comment='test_your_comment', filename_suffix="_test_your_filename_suffix")

    # img 1     random
    fake_img = torch.randn(3, 512, 512)
    writer.add_image("fake_img", fake_img, 1)
    time.sleep(1)

    # img 2     ones
    fake_img = torch.ones(3, 512, 512)
    time.sleep(1)
    writer.add_image("fake_img", fake_img, 2)

    # img 3     1.1
    fake_img = torch.ones(3, 512, 512) * 1.1
    time.sleep(1)
    writer.add_image("fake_img", fake_img, 3)

    # img 4     HW
    fake_img = torch.rand(512, 512)
    writer.add_image("fake_img", fake_img, 4, dataformats="HW")

    # img 5     HWC
    fake_img = torch.rand(512, 512, 3)
    writer.add_image("fake_img", fake_img, 5, dataformats="HWC")

    writer.close()


# ----------------------------------- 4 make_grid -----------------------------------
# flag = 0
flag = 1
if flag:
    writer = SummaryWriter(comment='test_your_comment', filename_suffix="_test_your_filename_suffix")

    split_dir = os.path.join("..", "..", "data", "rmb_split")
    train_dir = os.path.join(split_dir, "train")
    # train_dir = "S:/MyHazeset/256hazeset/train_haze_set/"

    transform_compose = transforms.Compose([transforms.Resize((256, 256)), transforms.ToTensor()])
    train_data = RMBDataset(data_dir=train_dir, transform=transform_compose)
    train_loader = DataLoader(dataset=train_data, batch_size=16, shuffle=True)
    data_batch, label_batch = next(iter(train_loader))

    img_grid = vutils.make_grid(data_batch, nrow=4, normalize=True, scale_each=True)
    # img_grid = vutils.make_grid(data_batch, nrow=4, normalize=False, scale_each=False)
    writer.add_image("input img", img_grid, 0)

    writer.close()


# ----------------------------------- 5 add_graph -----------------------------------

flag = 0
# flag = 1
if flag:

    writer = SummaryWriter(comment='test_your_comment', filename_suffix="_test_your_filename_suffix")

    # 模型
    fake_img = torch.randn(1, 3, 32, 32)

    lenet = LeNet(classes=2)

    writer.add_graph(lenet, fake_img)

    writer.close()

    from torchsummary import summary
    print(summary(lenet, (3, 32, 32), device="cpu"))

add_image()

add_image(tag, img_tensor, global_step=None, walltime=None, dataformats='CHW')

功能：记录图像

• tag: 图像的标签名，图的唯一标识
• img_tensor：图像数据，注意尺度。若像素值为0-1,则会默认乘以255；若像素值有大于1的，则默认为0-255
• global_step：x轴
• dataformats: 数据形式，CHW，HWC，HW

接下来我们运行上面代码的 3 image 部分，并用tensorboard打开。第一张由于使用的randn，所以为随机噪声图片；第二张初始化为1，程序默认乘以255，所以为全白图片；第三张像素值由于大于1，所以在0-255的区间所有像素值为1.1，图片为黑色；第四张为灰度图且为随机噪声图片；第五张为随机噪声图。

为了在一个界面中看到所有图片，此时我们引入torchvision.utils.make_grid。

torchvision.utils.make_grid

make_grid(tensor, nrow=8, padding=2, normalize=False, range=None, scale_each=False, pad_value=0)

功能：制作网格图像

• tensor：图像数据，BCH*W形式
• nrow: 行数（列数自动计算）
• padding：图像间距（像素单位）
• normalize: 是否将像素值标准化
• range：标准化范围
• scale_each: 是否单张图维度标准化
• pad_value: padding的像素值

我们运行 4 make_grid 部分代码，并用tensorboard打开，效果图如下。
True

False

add_graph()

add_graph(model, input_to_model=None, verbose=False)

功能：可视化模型计算图

• model: 模型，必须是nn.Module
• input_to_model: 输出给模型的参数
• verbose：是否打印计算图结构信息

运行 5 add_graph的代码，并用tensorboard打开。

卷积核可视化

import torch.nn as nn
from PIL import Image
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
import torchvision.utils as vutils
import os
import sys
hello_pytorch_DIR = os.path.abspath(os.path.dirname(__file__)+os.path.sep+".."+os.path.sep+"..")
sys.path.append(hello_pytorch_DIR)

from tools.common_tools import set_seed
import torchvision.models as models

set_seed(1)  # 设置随机种子


# ----------------------------------- kernel visualization -----------------------------------
# flag = 0
flag = 1
if flag:
    writer = SummaryWriter(comment='test_your_comment', filename_suffix="_test_your_filename_suffix") # 构建writer
	# 获取预训练的alexnet
    alexnet = models.alexnet(pretrained=True)

    kernel_num = -1 # 当前是第几个卷积层
    vis_max = 1  # 最大可视化层

    for sub_module in alexnet.modules():
        if isinstance(sub_module, nn.Conv2d):
            kernel_num += 1
            if kernel_num > vis_max:
                break
            kernels = sub_module.weight
            c_out, c_int, k_w, k_h = tuple(kernels.shape)

            for o_idx in range(c_out):
                kernel_idx = kernels[o_idx, :, :, :].unsqueeze(1)   # make_grid需要 BCHW，这里拓展C维度
                kernel_grid = vutils.make_grid(kernel_idx, normalize=True, scale_each=True, nrow=c_int)
                writer.add_image('{}_Convlayer_split_in_channel'.format(kernel_num), kernel_grid, global_step=o_idx)

            kernel_all = kernels.view(-1, 3, k_h, k_w)  # 3, h, w
            kernel_grid = vutils.make_grid(kernel_all, normalize=True, scale_each=True, nrow=8)  # c, h, w
            writer.add_image('{}_all'.format(kernel_num), kernel_grid, global_step=322)

            print("{}_convlayer shape:{}".format(kernel_num, tuple(kernels.shape)))

    writer.close()


# ----------------------------------- feature map visualization -----------------------------------
# flag = 0
flag = 1
if flag:
    writer = SummaryWriter(comment='test_your_comment', filename_suffix="_test_your_filename_suffix")

    # 数据
    path_img = "./lena.png"     # your path to image
    normMean = [0.49139968, 0.48215827, 0.44653124]
    normStd = [0.24703233, 0.24348505, 0.26158768]

    norm_transform = transforms.Normalize(normMean, normStd)
    img_transforms = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        norm_transform
    ])

    img_pil = Image.open(path_img).convert('RGB')
    if img_transforms is not None:
        img_tensor = img_transforms(img_pil)
    img_tensor.unsqueeze_(0)    # chw --> bchw

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

    # forward
    convlayer1 = alexnet.features[0]
    fmap_1 = convlayer1(img_tensor)

    # 预处理
    fmap_1.transpose_(0, 1)  # bchw=(1, 64, 55, 55) --> (64, 1, 55, 55)
    fmap_1_grid = vutils.make_grid(fmap_1, normalize=True, scale_each=True, nrow=8)

    writer.add_image('feature map in conv1', fmap_1_grid, global_step=322)
    writer.close()
    

代码结果如下

特征图可视化

torchsummary

summary(model, input_size, batch_size=-1, device="cuda")

功能：查看模型信息，便于调试

• model: pytorch模型
• input_size: 模型输入size
• batch_size: batch size
• device: “cuda” or “cpu”
  github: https://github.com/sksq96/pytorch-summary