pytorch 利用netron可视化网络结构

if __name__ == '__main__':
    import torch.onnx
    import netron

#####hourglass是自己的网络代码，自己定义的网络结构类名
    pose = HourglassNet(Bottleneck, 256, 8, 1, 16)  # .cuda()
    dummy_input = torch.randn(1, 3,256, 256)
#输出的文件名称，一般是在当前定义网络路径下
    onnx_path= "pose.onnx"
    torch.onnx.export(pose, dummy_input, "pose.onnx")  # netron --host=localhost
#自动跳转到netron的网址下
    netron.start(onnx_path)