保存方式一：

例如：對 vgg16 網絡模型進行保存，模型如下

模型保存使用 torch.save() 方法

torch.save(模型, "文件名.h5") # 保存 模型結構 + 模型參數

使用方式如下：

import torch import torchvisionvgg16 = torchvision.models.vgg16(pretrained=False) # 加載vgg16網絡模型# 保存方式1 torch.save(vgg16, "vgg16_method.pth") # 保存為 vgg16_method.pth 文件

加載保存后的模型文件

import torchmodel = torch.load("vgg16_method.pth") print(model)

保存方式二：

import torch import torchvisionvgg16 = torchvision.models.vgg16(pretrained=False) # 加載vgg16網絡模型# 保存模型的參數 torch.save(vgg16.state_dict(), "vgg16_method.pth") # 把模型的參數 保存成字典

加載保存后的模型文件

import torchmodel = torch.load("vgg16_method.pth") # 使用 torch.load 加載 模型文件，得到模型的參數 print(model)

使用以下方法得到完整的網絡模型

import torch import torchvisionvgg16 = torchvision.models.vgg16(pretrained=False) vgg16.load_state_dict(torch.load("vgg16_method.pth")) print(vgg16)

保存自定義模型：

下面是自定義了一個 VVcatModel 模型進行保存。

import torch from torch import nnclass VVcatModel(nn.Module):def __init__(self):super().__init__()self.conv1 = nn.Conv2d(3, 64, kernel_size=3)def forward(self, x):x = self.conv1(x)return xvvcat = VVcatModel() # 實例化模型 torch.save(vvcat, "vgg16_method.pth")

加載自定義模型

import torch from torch import nnclass VVcatModel(nn.Module):def __init__(self):super().__init__()self.conv1 = nn.Conv2d(3, 64, kernel_size=3)def forward(self, x):x = self.conv1(x)return x# 此處不需要實例化模型 model = torch.load("vgg16_method.pth") print(model)

注：自定義保存模型與加載自定義模型區別在于，加載模型時，不需要實例化模型。

注意：
如果使用GPU去訓練生成的網絡模型文件，在加載時可能出現以下錯誤

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

當使用GPU去訓練生成的網絡模型文件，在加載時使用的是CPU去加載，所以，在加載該網絡文件時，需要指定CPU去加載該文件， map_location=torch.device(‘cpu’)，添加如下：

model = torch.load("./vgg16_method.pth", map_location=torch.device('cpu'))

總結

以上是生活随笔為你收集整理的Pytorch 网络模型的保存与读取的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。