基于pytorch的保存和加載模型參數(shù)的方法

發(fā)布時(shí)間：2020-10-18 19:13:07 來源：腳本之家閱讀：175 作者：hit_zc 欄目：開發(fā)技術(shù)

當(dāng)我們花費(fèi)大量的精力訓(xùn)練完網(wǎng)絡(luò)，下次預(yù)測數(shù)據(jù)時(shí)不想再（有時(shí)也不必再）訓(xùn)練一次時(shí)，這時(shí)候torch.save(),torch.load()就要登場了。

保存和加載模型參數(shù)有兩種方式：

方式一：

torch.save(net.state_dict(),path):

功能：保存訓(xùn)練完的網(wǎng)絡(luò)的各層參數(shù)（即weights和bias)

其中：net.state_dict()獲取各層參數(shù)，path是文件存放路徑(通常保存文件格式為.pt或.pth)

net2.load_state_dict(torch.load(path)):

功能：加載保存到path中的各層參數(shù)到神經(jīng)網(wǎng)絡(luò)

注意：不可以直接為torch.load_state_dict(path)，此函數(shù)不能直接接收字符串類型參數(shù)

方式二：

torch.save(net,path):

功能：保存訓(xùn)練完的整個(gè)網(wǎng)絡(luò)模型（不止weights和bias）

net2=torch.load(path):

功能：加載保存到path中的整個(gè)神經(jīng)網(wǎng)絡(luò)

說明：官方推薦方式一，原因自然是保存的內(nèi)容少，速度會(huì)更快。

以上這篇基于pytorch的保存和加載模型參數(shù)的方法就是小編分享給大家的全部內(nèi)容了，希望能給大家一個(gè)參考，也希望大家多多支持億速云。

向AI問一下細(xì)節(jié)

猜你喜歡