Pytorch加载.pth文件

1. .pth文件

.pth文件报存了模型的权重，这个文件只是一个模型张量参数的pickle文件。

我们可以使用模型的load _state_dict方法将它们加载到 resnet18 中

2. 加载

2.1 如果.pth文件只保存了参数，则如下：

import torch
from torch.serialization import load
import torchvision.models as models

# pretrained=True使用预训练的模型
resnet18 = models.resnet18(pretrained=True)#创建实例，模型下载.Pth文件
model_path = 'D:/python_code/resnet18/resnet18-5c106cde.pth'
model_data = torch.load(model_path)
resnet18.load_state_dict(model_data)
print(resnet18)

输出为：

2.2 如果.pth文件保存的是整个网络结构+参数，则：

import torchvision.models as models
 
# pretrained=True就可以使用预训练的模型
resnet18 = models.resnet18(pretrained=True)
print(resnet18)

输出为：