使用hiddenlayer或torchviz可视化网络结构，并使用Tensorboard和wandb可视化训练过程

前言

本文参考极市平台的推文，出自知乎的一位锦恢的作者，使用并改写而来，链接如下：

PyTorch下的可视化工具（网络结构/训练过程可视化）

改写内容：将可视化方法应用在 ResNet 上，并展现出了灵活的可嵌入性。

仅作为笔记分享，如有侵权，请联系删除。

1. Graphviz

可视化神经网络的库 hiddenlayer 和 torchviz 都是基于 Graphviz 开发的， 使用前必须先安装

提示一下：Graphviz比较小，只有5M不到，放心安装，安装过程也比较简单

下载路径：https://www.graphviz.org/download/

我选择的是下图的红色框中的版本。

只需要注意一下，在安装过程中，添加到环境变量中，两个选一个就行了。

1.1 hiddenlayer、torchviz

安装好 Graphviz 后，直接使用命令行安装两个可视化库。

pip install hiddenlayer
pip install torchviz

这里先附上作者提供的代码，简单实用，看一下就明白了。

import torch
import torch.nn as nn

class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 16, 3, 1, 1),
            nn.ReLU(),
            nn.AvgPool2d(2, 2)
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )

        self.fc = nn.Sequential(
            nn.Linear(32 * 7 * 7, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU()
        )

        self.out = nn.Linear(64, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        output = self.out(x)
        return output

MyConvNet = ConvNet()
print(MyConvNet)

# method 1
import hiddenlayer as h

vis_graph = h.build_graph(MyConvNet, torch.zeros([1, 1, 28, 28]))  # 获取绘制图像的对象
vis_graph.theme = h.graph.THEMES["blue"].copy()  # 指定主题颜色
vis_graph.save("./demo1.png")  # 保存图像的路径

# method 2
# from torchviz import make_dot
#
# x = torch.randn(1, 1, 28, 28).requires_grad_(True)  # 定义一个网络的输入值
# y = MyConvNet(x)  # 获取网络的预测值
#
# MyConvNetVis = make_dot(y, params=dict(list(MyConvNet.named_parameters()) + [('x', x)]))
# MyConvNetVis.format = "png"
# # 指定文件生成的文件夹
# MyConvNetVis.directory = "data"
# # 生成文件
# MyConvNetVis.view()

代码解析：首先定义一个网络，之后直接使用 hiddenlaye 或 torchviz 即可，要注意给一个虚拟数据，维数和输入相同就行，具体详见上述作者链接。

需要改写到自己的网络中，只需要在网络模型实例化之后，加上这个操作即可。

使用 hiddenlayer 的可视化效果，我实例化的一个ResNet网络，比较长，只截取了一部分，缺点就是比较长。但结构显示效果还行。

同样是 ResNet，使用 torchviz 的可视化效果如下，总体感觉还行，长宽占比比较平衡，但就是分辨率不高，放大了看都是马赛克。目前还不清楚是什么原因，有知道了可以评论补充一下。

Tensorboard

使用Tensorboard可视化训练过程，其中最重要的就是使用 SummaryWriter 来记录下训练过程中你要可视化的数据。对的，就是比正常的网络多了这么一个操作。

如果没有TensorboardX库的，也需要先pip一下哈。

总的来说就是三部：

第一步：导入 TensorboardX 中的 SummaryWriter，并实例化一个logger，后面的地址是存放训练过程数据的地址。

logger = SummaryWriter(log_dir="data/log")

第二步：在网络训练过程中，使用 logger 记录训练过程数据，不会影响到网络中的其他部分，直接记录即可。代码如下：

同理，简单使用 ResNet 对CIFAR10 数据集进行训练。

多补充一句，以下的代码只需要pip install相关的库，然后在dataset部分，设置自己的数据保存路径，download=True即可。

import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import time as t
from tensorboardX import SummaryWriter

import wandb
logger = SummaryWriter(log_dir="data/log2")

# Device configuration
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Hyper-parameters
num_epochs = 5
batch_size = 100
learning_rate = 0.001

# Image preprocessing modules   预处理操作,模块化，预定义好一系列的预处理操作，在下面设置数据集的时候使用，用compose连接
transform = transforms.Compose([
    transforms.Pad(4),
    transforms.RandomHorizontalFlip(),
    transforms.RandomCrop(32),
    transforms.ToTensor()])

# CIFAR-10 dataset
train_dataset = torchvision.datasets.CIFAR10(root=r'D:\DATA\DL_Torch\DATAS\CIFAR10',
                                             train=True, 
                                             transform=transform,
                                             download=False)

test_dataset = torchvision.datasets.CIFAR10(root=r'D:\DATA\DL_Torch\DATAS\CIFAR10',
                                            train=False, 
                                            transform=transforms.ToTensor())

# Data loader
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=False)

# 3x3 convolution   定义一个基本的卷积模块
def conv3x3(in_channels, out_channels, stride=1):
    return nn.Conv2d(in_channels, out_channels, kernel_size=3, 
                     stride=stride, padding=1, bias=False)

# Residual block    定义一个基本的残差模块
class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(ResidualBlock, self).__init__()
        self.conv1 = conv3x3(in_channels, out_channels, stride)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(out_channels, out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = downsample
        
    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.downsample:
            residual = self.downsample(x)
        out += residual
        out = self.relu(out)
        return out

# ResNet
class ResNet(nn.Module):
    def __init__(self, block, layers, num_classes=10):
        super(ResNet, self).__init__()
        self.in_channels = 16
        self.conv = conv3x3(in_channels=3, out_channels=16)
        self.bn = nn.BatchNorm2d(16)
        self.relu = nn.ReLU(inplace=True)
        self.layer1 = self.make_layer(block, 16, layers[0])
        self.layer2 = self.make_layer(block, 32, layers[1], 2)
        self.layer3 = self.make_layer(block, 64, layers[2], 2)
        self.avg_pool = nn.AvgPool2d(8)
        self.fc = nn.Linear(64, num_classes)
        
    def make_layer(self, block, out_channels, blocks, stride=1):
        downsample = None
        if (stride != 1) or (self.in_channels != out_channels):
            downsample = nn.Sequential(
                conv3x3(self.in_channels, out_channels, stride=stride),
                nn.BatchNorm2d(out_channels))
        layers = []
        layers.append(block(self.in_channels, out_channels, stride, downsample))
        self.in_channels = out_channels
        for i in range(1, blocks):
            layers.append(block(out_channels, out_channels))
        return nn.Sequential(*layers)
    
    def forward(self, x):
        out = self.conv(x)
        out = self.bn(out)
        out = self.relu(out)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.avg_pool(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out

# ----------------------------------- log set ---------------------------------------------
# 初始化一个wandb run
wandb.init(project="Res-Net")
# ----------------------------------- model ---------------------------------------------
model = ResNet(ResidualBlock, [2, 2, 2])

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# For updating learning rate  实时更新学习率函数
def update_lr(optimizer, lr):
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

# Train the unet_model，标准训练，加了个学习率衰减
total_step = len(train_loader)
curr_lr = learning_rate
start = t.time()
# ----------------------------------- train ---------------------------------------------
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # images = images.to(device)
        # labels = labels.to(device)

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)

        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        global_iter_num = epoch * len(train_loader) + i + 1

        if (i+1) % 100 == 0:
            print ("Epoch [{}/{}], Step [{}/{}] Loss: {:.4f}"
                   .format(epoch+1, num_epochs, i+1, total_step, loss.item()))
            logger.add_scalar("train loss", loss.item(), global_step=global_iter_num)


            logger.add_image("train image sample", images[-3], global_step=global_iter_num)

            for name, param in model.named_parameters():
                logger.add_histogram(name, param.data.numpy(), global_step=global_iter_num)

            # 定义需要记录的数据
            wandb.log({
                'train_image_sample' : images[-3],
                'train_loss' : loss.item(),
            })

    # Decay learning rate
    if (epoch+1) % 20 == 0:
        curr_lr /= 3
        update_lr(optimizer, curr_lr)

# Test the unet_model,标准模式测试
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for j, images, labels in enumerate(test_loader):
        # images = images.to(device)
        # labels = labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        global_test = len(test_loader) + j + 1

    print('Accuracy of the unet_model on the test images: {} %'.format(100 * correct / total))
    logger.add_scalar("test accuary", correct, global_step=global_test)
    # 使用wandb.log()记录需要保存的信息
    wandb.log({
        'predict_image' : outputs
    })

# 记录一下模型中的所有参数
wandb.watch(model, log="all")
# Save the unet_model checkpoint
torch.save(model.state_dict(), 'resnet_gpu.ckpt')
end = t.time()
print(end-start)

上述代码简单记录了一下ResNet的训练过程，代码应该比较易读，没有什么复杂的结构化的东西。

在训练结束之后，可以在之前设定存放logger数据目录的地方找到一个文件。

方法一：使用Terminal命令行：tensorboard --logdir=./data/log，后面的是log的地址，可以在pycharm的Terminal中输入，会得到一个地址，单击打开即可，或复制到游览器中打开。

方法二：使用终端命令打开

进入cmd -->>-- 激活conda环境 -->>-- 输入tensorboard --logdir=logdir/path。完成

注意：这里的地址地方不需要加上引号。

最后打开，就可以看到可视化效果了，总体来说，体验还是很好的。

Tensorboard概括

简短的说明一下tensorboard怎么加进网络中。

# 1.导入库，并且实例化一个logger，设置保存的位置
from tensorboardX import SummaryWriter
logger = SummaryWriter(log_dir="data/log2")

# 2. 加载数据集、模型、优化器、损失函数等
model = Net()

# 3. 开始训练，记录中间数据，scalar---记录点，image--记录图片
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
    	# ---------------------------
		logger.add_scalar("train loss", loss.item(), global_step=global_iter_num)
		logger.add_image("train image sample", images[-3], global_step=global_iter_num)
		logger.add_image("train image sample", make_grid(out[0].detach().cpu().unsqueeze(dim=1), nrow=3, padding=20, normalize=False, pad_value=255), global_step=1)		# 可视化所有通道的特征图
		logger.add_graph(model, x)	# logger.add_graph(model, [x, y])

# 4. 模型验证
model.eval()
with torch.no_grad():
	for j, images, labels in enumerate(test_loader):
		# ---------------------------
		logger.add_scalar("test accuary", correct, global_step=global_test)

可以看出，就是在想要记录的位置加入相关的命令即可。

其他的自己尝试一下感受一下吧！

Tensorboard实例化效果

其他的一些操作可以参考官方文档：

https://pytorch.org/docs/stable/search.html?q=tensorboard&check_keywords=yes&area=default#

wandb

使用wandb可以参考官方的文档，Weights & Biases

详细写了一些简单的用法和示例，别没有耐心看，还是比较人性化的。

按照他的来就行，需要注册一下，获取自己的key，都按照文档提示操作即可。

上述的代码中也加入了wandb可视化库，下面直接就使用进行概括

# 导入库
import wandb
# 初始化一个wandb，名字可以自己定义
wandb.init(project="Res-Net")

# 2. 加载数据集、模型、优化器、损失函数等

# 3. 开始训练，记录中间数据，直接使用log方法
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
    	# ---------------------------
    	# 定义需要记录的数据
            wandb.log({
                'train_image_sample' : images[-3],
                'train_loss' : loss.item(),
            })

# 4. 模型验证
model.eval()
with torch.no_grad():
	for j, images, labels in enumerate(test_loader):
		# ---------------------------
		wandb.log({
        'predict_image' : outputs
    })

# 5. 记录一下模型中的所有参数
wandb.watch(model, log="all")

可以看到，wandb使用方法和tensorboard类似，都比较简单。

wandb实例化效果：

这个可以直接在网络中增加，很方便，不需要过多的改动，效果也很好！赞！！！

这里我个人还是比较喜欢tensorboard，但是wandb可以很方便的查看项目的记录信息，而且可以看到计算机CPU和GPU的相关计算信息，打开也很方便。

使用wandb在最后的保存上会占用一点时间，可能是保存的信息比较多，相对比较慢。

根据自己的喜好选择吧。

生成的文件

使用tensorboard和wandb生成的文件

最后感谢锦恢，在此仅作笔记分享，如有侵权，联系删除。