pytorch模型結構可視化，可顯示每層的尺寸
- 2021 年 1 月 10 日
- 筆記
- Deep Learning, 深度學習, PyTorch
最近在學習一些檢測方面的網路，使用的是pytorch。模型結構可視化是學習網路的有用的部分，pytorch沒有原生支援這個功能，需要找一些其他方式，下面總結幾種方法（推薦用4）。
1. torch .pt -> netron
netron是一個專門可視化模型的工具，支援很多格式，很自然的就是用它直接顯示torch保存的模型。但是實際上...
顯示成了上圖，基本上沒什麼用。
2. 網上說的比較多的幾種方式
a. tensorboardx 畫出來的比較丑
b. tensorwatch 支援的torch版本低，容易出錯
c. graphviz 實測yolov5異常。。
3. torch->caffe->netron
caffe支援的操作不夠多，轉換很容易碰到不支援的情況
4. torch->onnx->netron
pytorch對onnx支援的比較好，直接有導出介面。
torch.onnx.export...
但這樣直接轉出來的onnx進行顯示沒有每一層的輸入大小，如下圖，還是不方便看。
我們可以將導出的onnx進行以下操作：
import onnx
from onnx import shape_inference
model = 'best.onnx'
onnx.save(onnx.shape_inference.infer_shapes(onnx.load(model)), model)
再用netron進行顯示即可得到，用這種結構圖學習就好多了：
歡迎關註：老司機的視覺屋
