pytorch中模型和参数的存储

最新推荐文章于 2024-08-29 22:12:23 发布

原创最新推荐文章于 2024-08-29 22:12:23 发布 · 1k 阅读

1 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习

本文详细介绍了PyTorch中torch.nn.Module模型的state_dict属性，该属性是Python字典对象，用于存储模型的可学习参数，如权重和偏差。同时，也讲解了torch.optim优化器的state_dict属性，包含了优化器的状态信息和超参数。文章通过一个分类器模型实例，演示了如何访问和打印模型及优化器的state_dict。

部署运行你感兴趣的模型镜像

在PyTorch中， torch.nn.Module 模型的可学习参数（即权重和偏差）包含在模型的参数中，
（使用 c可以进行访问）。
state_dict 是Python字典对象，它将每一层映射到其参数张量。注意，只有具有可学习参数的层（如卷积层，线性层等）的模型才具有 state_dict 这一项。目标优化 torch.optim 也有 state_dict 属性，它包含有关优化器的状态信息，以及使用的超参数。
因为state_dict的对象是Python字典，所以它们可以很容易的保存、更新、修改和恢复，为PyTorch模型和优化器添加了大量模块。
下面通过从简单模型训练一个分类器中来了解一下 state_dict 的使用

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
class theModelClass(nn.Module):
    def __init__(self):
        super(theModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(6,16, 5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16*5*5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


model = theModelClass()

# 初始化 优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 打印模型的状态字典
print("Model's state_dict:" )
for param_tensor in model.state_dict():
    print(param_tensor, '\t', model.state_dict()[param_tensor].size())

# 打印优化器的转台字典
print(" Optimizer's state_dict; ")
for var_name in optimizer.state_dict():
    print(var_name, '\t', optimizer.state_dict()[var_name])

# 输出

"""
Model's state_dict:
conv1.weight torch.Size([6, 3, 5, 5])
conv1.bias torch.Size([6])

conv2.weight torch.Size([16, 6, 5, 5])
conv2.bias torch.Size([16])

fc1.weight torch.Size([120, 400])
fc1.bias torch.Size([120])

fc2.weight torch.Size([84, 120])
fc2.bias torch.Size([84])

fc3.weight torch.Size([10, 84])
fc3.bias torch.Size([10])

Optimizer's state_dict:
state {}
param_groups [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay':
0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072,
4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]