Pytorch-day07-模型保存与读取

PyTorch 模型保存&读取

• 模型存储
• 模型单卡存储&多卡存储
• 模型单卡读取&多卡读取

1、模型存储

• PyTorch存储模型主要采用pkl，pt，pth三种格式,就使用层面来说没有区别
• PyTorch模型主要包含两个部分：模型结构和权重。其中模型是继承nn.Module的类，权重的数据结构是一个字典（key是层名，value是权重向量）
• 存储也由此分为两种形式：存储整个模型（包括结构和权重）和只存储模型权重(推荐）。

import torch
from torchvision import models
model = models.resnet50(pretrained