【Pytorch】模型摘要信息获取、模型参数获取及模型保存的三种方法
目录
- 问题一：模型摘要信息的获取
- 问题二：模型参数的获取
- 问题三：模型的保存方式
问题1：我想得到模型的摘要信息，包括每一层的名称、输入尺寸、输出尺寸以及参数量。
PyTorch Summary是一个用于计算模型参数量和输出尺寸的工具库。它可以帮助你快速了解模型的结构和参数数量，以及每个层的输出形状。你可以使用torchsummary库来生成模型的摘要信息。以下是一个示例代码：
import torch
from torchvision import models
from torchsummary import summary
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg = models.vgg16().to(device)
summary(vgg, (3, 224, 224))
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 224, 224] 1,792
ReLU-2 [-1, 64, 224, 224] 0
Conv2d-3 [-1, 64, 224, 224] 36,928
ReLU-4 [-1, 64, 224, 224] 0
MaxPool2d-5 [-1, 64, 112, 112] 0
Conv2d-6 [-1, 128, 112, 112] 73,856
ReLU-7 [-1, 128, 112, 112] 0
Conv2d-8 [-1, 128, 112, 112] 147,584
ReLU-9 [-1, 128, 112, 112] 0
MaxPool2d-10 [-1, 128, 56, 56] 0
Conv2d-11 [-1, 256, 56, 56] 295,168
ReLU-12 [-1, 256, 56, 56] 0
Conv2d-13 [-1, 256, 56, 56] 590,080
ReLU-14 [-1, 256, 56, 56] 0
Conv2d-15 [-1, 256, 56, 56] 590,080
ReLU-16 [-1, 256, 56, 56] 0
MaxPool2d-17 [-1, 256, 28, 28] 0
Conv2d-18 [-1, 512, 28, 28] 1,180,160
ReLU-19 [-1, 512, 28, 28] 0
Conv2d-20 [-1, 512, 28, 28] 2,359,808
ReLU-21 [-1, 512, 28, 28] 0
Conv2d-22 [-1, 512, 28, 28] 2,359,808
ReLU-23 [-1, 512, 28, 28] 0
MaxPool2d-24 [-1, 512, 14, 14] 0
Conv2d-25 [-1, 512, 14, 14] 2,359,808
ReLU-26 [-1, 512, 14, 14] 0
Conv2d-27 [-1, 512, 14, 14] 2,359,808
ReLU-28 [-1, 512, 14, 14] 0
Conv2d-29 [-1, 512, 14, 14] 2,359,808
ReLU-30 [-1, 512, 14, 14] 0
MaxPool2d-31 [-1, 512, 7, 7] 0
Linear-32 [-1, 4096] 102,764,544
ReLU-33 [-1, 4096] 0
Dropout-34 [-1, 4096] 0
Linear-35 [-1, 4096] 16,781,312
ReLU-36 [-1, 4096] 0
Dropout-37 [-1, 4096] 0
Linear-38 [-1, 1000] 4,097,000
================================================================
Total params: 138,357,544
Trainable params: 138,357,544
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 218.59
Params size (MB): 527.79
Estimated Total Size (MB): 746.96
----------------------------------------------------------------
问题2：model.parameters()与model.state_dict()是干嘛的？
-
model.parameters(): 这个方法返回一个包含模型所有可学习参数的迭代器。可学习参数包括模型的权重(weights)和偏置(biases)等需要通过梯度更新的参数。model.parameters()常用于定义优化器(optimizer)和计算梯度。
-
model.state_dict(): 这个方法返回一个字典，包含了模型的所有状态信息。字典中的键是参数名称，值是对应参数的张量(Tensor)。model.state_dict()的主要用途是保存和加载模型。通过调用torch.save()将model.state_dict()保存为文件后，可以使用torch.load()加载模型参数并将其应用到模型中。
import torch
import torch.nn as nn
# 定义一个简单的模型
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc1 = nn.Linear(5, 2)
def forward(self, x):
x = self.fc1(x)
return x
# 创建模型实例
model = MyModel()
# 打印模型的可学习参数
for param in model.parameters():
print(param, param.shape)
print (model.state_dict())
Parameter containing:
tensor([[ 4.1945e-01, 3.8990e-01, 3.7970e-01, -1.8200e-04, 3.4936e-01],
[-9.5073e-02, 8.0670e-02, -2.4634e-01, -3.7250e-01, 2.4676e-01]],
requires_grad=True) torch.Size([2, 5])
Parameter containing:
tensor([ 0.3537, -0.2398], requires_grad=True) torch.Size([2])
----------
OrderedDict([('fc1.weight', tensor([[ 4.1945e-01, 3.8990e-01, 3.7970e-01, -1.8200e-04, 3.4936e-01],
[-9.5073e-02, 8.0670e-02, -2.4634e-01, -3.7250e-01, 2.4676e-01]])), ('fc1.bias', tensor([ 0.3537, -0.2398]))])
问题3：Pytorch模型保存的几种方法？
模型保存的方式取决于你后续加载模型的用途。
- 保存模型以供自己用于推理：保存模型，恢复模型，然后将模型更改为评估模式。
torch.save(model.state_dict(), filepath)
#Later to restore:
model.load_state_dict(torch.load(filepath))
model.eval()
- 保存模型以便稍后恢复训练：如果需要继续训练你将要保存的模型，那么需要保存的不仅仅是模型。还需要保存优化器的状态、迭代次数、评估指标等。可以这样做：
state = {
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
...
}
torch.save(state, filepath)
# 要恢复训练，并恢复每个单独对象的状态，如下所示：
state = torch.load(filepath)
model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])
- 模型被无法访问你代码的其他人使用：
torch.save(model, filepath)
# Then later:
model = torch.load(filepath)
请关注博主，一起玩转人工智能及深度学习。