目的:
把訓(xùn)練好的pth模型參數(shù)提取出來(lái),然后用其他方式部署到邊緣設(shè)備。
Pytorch給了很方便的讀取參數(shù)接口:
直接看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參數(shù)了~
Note:
model.parameters()是以一個(gè)生成器的形式迭代返回每一層的參數(shù)。所以用for循環(huán)讀取到各層的參數(shù),循環(huán)次數(shù)就表示層數(shù)。
而每一層的參數(shù)都是torch.nn.parameter.Parameter類型,是Tensor的子類,所以直接用tensor轉(zhuǎn)numpy(即p.detach().cpu().numpy())的方法就可以直接轉(zhuǎn)成numpy矩陣。
方便又好用,爆贊~
補(bǔ)充:pytorch訓(xùn)練好的.pth模型轉(zhuǎn)換為.pt
將python訓(xùn)練好的.pth文件轉(zhuǎn)為.pt
import torch
import torchvision
from unet import UNet
model = UNet(3, 2)#自己定義的網(wǎng)絡(luò)模型
model.load_state_dict(torch.load("best_weights.pth"))#保存的訓(xùn)練模型
model.eval()#切換到eval()
example = torch.rand(1, 3, 320, 480)#生成一個(gè)隨機(jī)輸入維度的輸入
traced_script_module = torch.jit.trace(model, example)
traced_script_module.save("model.pt")
以上為個(gè)人經(jīng)驗(yàn),希望能給大家一個(gè)參考,也希望大家多多支持腳本之家。如有錯(cuò)誤或未考慮完全的地方,望不吝賜教。
您可能感興趣的文章:- Numpy實(shí)現(xiàn)矩陣運(yùn)算及線性代數(shù)應(yīng)用
- numpy數(shù)組合并和矩陣拼接的實(shí)現(xiàn)
- numpy和tensorflow中的各種乘法(點(diǎn)乘和矩陣乘)
- NumPy 矩陣乘法的實(shí)現(xiàn)示例
- Python numpy大矩陣運(yùn)算內(nèi)存不足如何解決
- 使用numpy實(shí)現(xiàn)矩陣的翻轉(zhuǎn)(flip)與旋轉(zhuǎn)