天天看點

pytorch模型結構可視化,可顯示每層的尺寸

最近在學習一些檢測方面的網絡,使用的是pytorch。模型結構可視化是學習網絡的有用的部分,pytorch沒有原生支援這個功能,需要找一些其他方式,下面總結幾種方法(推薦用4)。

1. torch .pt -> netron

netron是一個專門可視化模型的工具,支援很多格式,很自然的就是用它直接顯示torch儲存的模型。但是實際上...

pytorch模型結構可視化,可顯示每層的尺寸

顯示成了上圖,基本上沒什麼用。

2. 網上說的比較多的幾種方式

a. tensorboardx 畫出來的比較醜

b. tensorwatch 支援的torch版本低,容易出錯

c. graphviz 實測yolov5異常。。

3. torch->caffe->netron

caffe支援的操作不夠多,轉換很容易碰到不支援的情況

4. torch->onnx->netron

pytorch對onnx支援的比較好,直接有導出接口。

torch.onnx.export...

但這樣直接轉出來的onnx進行顯示沒有每一層的輸入大小,如下圖,還是不友善看。

pytorch模型結構可視化,可顯示每層的尺寸

我們可以将導出的onnx進行以下操作:

import onnx

from onnx import shape_inference

model = \'best.onnx\'

onnx.save(onnx.shape_inference.infer_shapes(onnx.load(model)), model)

再用netron進行顯示即可得到,用這種結構圖學習就好多了:

pytorch模型結構可視化,可顯示每層的尺寸

 歡迎關注:老司機的視覺屋