给出完整的训练、保存模型,加载使用模型的例子
时间: 2024-01-21 11:16:34 浏览: 20
好的,我可以给你一个使用PyTorch实现的例子。我们将使用一个简单的神经网络模型来训练MNIST手写数字数据集,并保存和加载模型。
首先,让我们导入所需的库和模块:
```python
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
```
然后,定义一个简单的神经网络模型。这个模型由两个全连接层组成,使用ReLU激活函数,并在最后一层使用softmax函数输出类别概率。
```python
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(784, 256)
self.fc2 = nn.Linear(256, 10)
def forward(self, x):
x = x.view(-1, 784)
x = nn.functional.relu(self.fc1(x))
x = nn.functional.softmax(self.fc2(x), dim=1)
return x
```
接下来,加载MNIST数据集并定义数据转换:
```python
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST(
root='./data', train=True, download=True, transform=transform
)
test_dataset = datasets.MNIST(
root='./data', train=False, download=True, transform=transform
)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=True)
```
然后,定义优化器和损失函数:
```python
model = Net()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
```
接下来,训练模型:
```python
epochs = 5
for epoch in range(epochs):
for i, (images, labels) in enumerate(train_loader):
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
if i % 100 == 0:
print(f'Epoch {epoch}, Step {i}, Loss {loss.item():.4f}')
```
最后,保存模型:
```python
torch.save(model.state_dict(), 'mnist_model.pt')
```
现在,我们已经成功地训练并保存了模型。接下来,我们将展示如何加载模型并使用它进行推理。
首先,加载模型:
```python
model = Net()
model.load_state_dict(torch.load('mnist_model.pt'))
```
然后,使用模型进行推理:
```python
model.eval()
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Accuracy: {correct / total:.4f}')
```
这就是一个完整的训练、保存和加载模型的例子。