pytorch利用tensorboardx可视化网络结构图

1、对于自己网络有外部参数需要输入时：如下所示，外部参数有：nstack, inp_dim, oup_dim ，bn

class PoseNet(nn.Module):
    def __init__(self, nstack, inp_dim, oup_dim, bn=False, increase=128, init_weights=True, **kwargs):

 2、定义可视化结构

       1）首先，对网络PoseNet传入参数输入：如，nstack=4； inp_dim=256； oup_dim =54；bn = True

       2）其次，定义输入。input =torch.rand(1,128,128,3)

      3)最后，在如下代码的comment 处写上pose

with SummaryWriter(comment='pose') as w:
        w.add_graph(pose, (input,))

if __name__ == '__main__':
    from time import time

    pose = PoseNet(4, 256, 54, bn=True)  # .cuda()
    for param in pose.parameters():
        if param.requires_grad:
            print('param autograd')
            break

    t0 = time()
    input = torch.rand(1, 128, 128, 3)  # .cuda()
    print(pose)
    output = pose(input)  # type: torch.Tensor
    output[0][0].sum().backward()

    with SummaryWriter(comment='pose') as w:
        w.add_graph(pose, (input,))

    t1 = time()