從Pytorch模型pth文件中讀取參數成numpy矩陣的操作

Posted on by

目的：

把訓練好的pth模型參數提取出來，然後用其他方式部署到邊緣設備。

Pytorch給瞭很方便的讀取參數接口：

nn.Module.parameters()

直接看demo：

from torchvision.models.alexnet import alexnet 
model = alexnet(pretrained=True).eval().cuda()
parameters = model.parameters()
for p in parameters:
  numpy_para = p.detach().cpu().numpy()
  print(type(numpy_para))
  print(numpy_para.shape)

上面得到的numpy_para就是numpy參數瞭~

Note：

model.parameters()是以一個生成器的形式迭代返回每一層的參數。所以用for循環讀取到各層的參數，循環次數就表示層數。

而每一層的參數都是torch.nn.parameter.Parameter類型，是Tensor的子類，所以直接用tensor轉numpy（即p.detach().cpu().numpy()）的方法就可以直接轉成numpy矩陣。

方便又好用，爆贊~

補充：pytorch訓練好的.pth模型轉換為.pt

將python訓練好的.pth文件轉為.pt

import torch
import torchvision
from unet import UNet
model = UNet(3, 2)#自己定義的網絡模型
model.load_state_dict(torch.load("best_weights.pth"))#保存的訓練模型
model.eval()#切換到eval（）
example = torch.rand(1, 3, 320, 480)#生成一個隨機輸入維度的輸入
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("model.pt")

以上為個人經驗，希望能給大傢一個參考，也希望大傢多多支持WalkonNet。如有錯誤或未考慮完全的地方，望不吝賜教。

推薦閱讀：