【Pytorch】模型摘要信息获取、模型参数获取及模型保存的三种方法

2023-10-17 16:28:15 浏览数 (1)

问题1:我想得到模型的摘要信息,包括每一层的名称、输入尺寸、输出尺寸以及参数量。

PyTorch Summary是一个用于计算模型参数量和输出尺寸的工具库。它可以帮助你快速了解模型的结构和参数数量,以及每个层的输出形状。你可以使用torchsummary库来生成模型的摘要信息。以下是一个示例代码:

代码语言:javascript复制
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))
代码语言:javascript复制
----------------------------------------------------------------
        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()是干嘛的?

  1. model.parameters(): 这个方法返回一个包含模型所有可学习参数的迭代器。可学习参数包括模型的权重(weights)和偏置(biases)等需要通过梯度更新的参数。model.parameters()常用于定义优化器(optimizer)和计算梯度。
  2. model.state_dict(): 这个方法返回一个字典,包含了模型的所有状态信息。字典中的键是参数名称,值是对应参数的张量(Tensor)。model.state_dict()的主要用途是保存和加载模型。通过调用torch.save()将model.state_dict()保存为文件后,可以使用torch.load()加载模型参数并将其应用到模型中。
代码语言:javascript复制
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())
代码语言:javascript复制
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模型保存的几种方法?

模型保存的方式取决于你后续加载模型的用途。

  1. 保存模型以供自己用于推理:保存模型,恢复模型,然后将模型更改为评估模式。
代码语言:javascript复制
torch.save(model.state_dict(), filepath)

#Later to restore:
model.load_state_dict(torch.load(filepath))
model.eval()
  1. 保存模型以便稍后恢复训练:如果需要继续训练你将要保存的模型,那么需要保存的不仅仅是模型。还需要保存优化器的状态、迭代次数、评估指标等。可以这样做:
代码语言:javascript复制
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'])
  1. 模型被无法访问你代码的其他人使用:
代码语言:javascript复制
torch.save(model, filepath)

# Then later:
model = torch.load(filepath)

0 人点赞