PyTorch状态字典:state_dict使用详解

PyTorch中的state_dict是一个python字典对象，将每个层映射到其参数Tensor。state_dict对象存储模型的可学习参数，即权重和偏差，并且可以非常容易地序列化和保存。在本篇文章中，我们将详细介绍PyTorch中的state_dict对象及其使用方法。

保存模型和state_dict

首先，我们来看如何将模型的state_dict保存到文件中。我们可以使用torch.save函数实现。例如，对于一个简单的神经网络模型，我们可以这样保存它的state_dict:

import torch
import torch.nn as nn

# 定义一个模型
class Net(nn.Module):
    def __init__(self):
        super(Net, 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

# 初始化一个模型
net = Net()

# 保存模型的state_dict
torch.save(net.state_dict(), 'model_state_dict.pth')

加载模型和state_dict

接下来，我们来看如何加载模型的state_dict。同样，我们可以使用torch.load函数。需要注意的是，在加载模型之前必须先实例化模型。这是因为模型结构需要匹配，否则会出现参数维度不一致的问题。

# 实例化一个Net模型
net = Net()

# 加载之前保存的state_dict
net.load_state_dict(torch.load('model_state_dict.pth'))

通过这个简单的例子，我们可以了解如何保存和加载模型的state_dict对象。在实际应用中，我们通常需要保存训练过程中的模型状态。接下来，我们将通过一个示例来演示如何保存和加载训练过程中的模型状态。

保存和加载训练过程中的模型状态

在训练过程中，我们通常会采用epoch作为单位来保存模型的状态。这样，我们就可以在训练完成后再次加载模型，并从上一个epoch继续训练。下面是一个保存和加载训练过程中模型状态的示例：

# 定义一个模型
class Net(nn.Module):
    def __init__(self):
        super(Net, 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

# 定义训练函数
def train(model, optimizer, loss_func, trainloader):
    for epoch in range(10):
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data

            optimizer.zero_grad()

            outputs = model(inputs)
            loss = loss_func(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if i % 2000 == 1999:
                print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
                running_loss = 0.0

        # 保存每个epoch之后的模型参数
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss
        }, 'model_epoch_{}.pth'.format(epoch))

# 实例化一个模型
net = Net()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 加载训练数据
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)

# 训练模型并保存每个epoch之后的模型状态
train(net, optimizer, criterion, trainloader)

# 加载最后一个epoch的模型状态
checkpoint = torch.load('model_epoch_9.pth')
net.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
loss = checkpoint['loss']
epoch = checkpoint['epoch']

# 继续训练
for epoch in range(10):
    pass

通过这个示例，我们可以看到如何保存和加载训练过程中的模型状态。在实际应用中，我们可以使用PyTorch提供的自动化工具（如torch.utils.data.DataLoader）和训练循环（如torch.optim.SGD）来构建更加复杂的训练过程，并保存训练过程中的模型状态。