PyTorch模型的保存与加载方法实例

以下是PyTorch模型的保存与加载方法实例的详细攻略：

PyTorch提供了多种方法来保存和加载模型，包括使用pickle、torch.save和torch.load等方法。以下是使用torch.save和torch.load方法保存和加载模型的详细步骤：

1. 定义模型并训练模型。

```python
import torch
import torch.nn as nn
import torch.optim as optim

# 定义模型
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

# 定义数据集和数据加载器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)

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

# 训练模型
for epoch in range(10):
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
inputs, labels = data
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(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

print('Finished Training')
```

1. 保存模型。

python
# 保存模型
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

这个代码会将模型的参数保存到指定的文件中。

1. 加载模型。

python
# 加载模型
net = Net()
net.load_state_dict(torch.load(PATH))

这个代码会从指定的文件中加载模型的参数，并将其应用到新的模型中。

以下是两个示例说明：

示例1：使用保存的模型进行预测

以下是一个使用保存的模型进行预测的示例代码：

import torch
import torchvision.transforms as transforms
from PIL import Image

# 定义数据预处理
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# 加载模型
net = Net()
net.load_state_dict(torch.load(PATH))

# 加载图片并进行预测
image = Image.open('test.jpg')
image = transform(image)
image = image.unsqueeze(0)
output = net(image)
_, predicted = torch.max(output, 1)
print(predicted)

在这个示例中，我们首先定义了数据预处理方式，然后使用Net类加载模型，并使用load_state_dict方法从文件中加载模型的参数。接着，我们使用PIL库加载图片，并进行数据处理。最后，我们使用训练好的模型对图片进行预测，并输出预测结果。

示例2：使用保存的模型进行微调

以下是一个使用保存的模型进行微调的示例代码：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.datasets as datasets

# 定义数据预处理
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# 加载数据集
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

# 定义数据加载器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)

# 加载模型
net = Net()
net.load_state_dict(torch.load(PATH))

# 将模型的最后一层替换为新的全连接层
num_ftrs = net.fc3.in_features
net.fc3 = nn.Linear(num_ftrs, 2)

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

# 训练模型
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(train_loader, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(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

print('Finished Training')

在这个示例中，我们首先定义了数据预处理方式，然后使用CIFAR10类加载数据集。接着，我们使用Net类加载模型，并使用load_state_dict方法从文件中加载模型的参数。我们将模型的最后一层替换为新的全连接层，并使用交叉熵损失函数和随机梯度下降优化器来微调模型。最后，我们使用微调后的模型对测试集进行预测，并输出预测准确率。