YOLOv5 完整网络模型可视化 pytorch netron

本文分享使用pytorch可视化工具netron查看YOLOv5网络模型结构。

正常在控制台pip命令安装

pip install netron

执行netron

C:\Users>netron
Serving at http://localhost:8080

网页打开

在YOLOv5中，打开网络模型主文件yolo.py。

将模型保存为一个pt文件

# Create model
model = Model(opt.cfg).to(device)
torch.save(model, "m.pt")

网页点击Open，或者直接把pt文件拖入，就能看到效果了。

但这一般不是我们想要的，因为子结构的细节都没有显示，所以有些同学觉得这个工具有点鸡肋，其实是有办法的：创建输入x，使用torch.jit.trace()来跟踪网络。

# Create model
model = Model(opt.cfg).to(device)
x = torch.randn(1, 3, 640, 640).to(device)
script_model = torch.jit.trace(model, x)
script_model.save("m.pt")

这样就可以看到完整网络结构了，全图太大，不往上扔了，同学们自己去看叭~