如何使用tensorboard解读网络模型

前言：

博主最近在搭建网络的时候，先搞清楚到底哪些参数进行了训练，网络数据流是不是和自己的预期一样。

更新日志：2022.09.16 1更

一、简单的实例-普通的nn.Sequential可视化

import math
import torch.nn as nn
import torch
from torch.nn import functional as F
from torchsummary import summary
from torchvision import models


def BasicBlock(in_ch,out_ch,stride=1):
    return nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, stride, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),  # inplace = True原地操作,节省显存
            nn.Conv2d(out_ch, out_ch, 3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )

if __name__ == "__main__":
    import imageio
    import numpy as np
    from torchviz import make_dot
    from torch.utils.tensorboard import SummaryWriter

    im_data_tensor = torch.random([3,512,512])
    model_t = BasicBlock(in_ch=3, out_ch=32)
    out = model_t(im_data_tensor)

    # 1. 来用tensorflow进行可视化
    with SummaryWriter("./log", comment="sample_model_visualization") as sw:
        sw.add_graph(model_t, im_data_tensor)

 应该不难看出其中的逻辑关系，每一个操作后面都有一个标号，

二、网络的结构和forward息息相关，

并行

    def forward(self, x):
        x0 = self.conv_br1(x)
        x1 = self.conv_br2(x)
        return x0,x1

 串行

    def forward(self, x):
        x0 = self.conv_br1(x)
        x1 = self.conv_br2(x0)
        return x1

 与最终​​​​​​​的输出无关的变量或者计算流，不会出现在网络中

    def forward(self, x):
        x0 = self.conv_br1(x)
        x1 = self.conv_br2(x)
        return x1