使用PyTorch训练MNIST数据集上的卷积神经网络
时间: 2024-12-22 13:15:02 浏览: 13
在PyTorch中,训练MNIST数据集上的卷积神经网络(Convolutional Neural Network,CNN)通常包括以下几个步骤:
1. **导入必要的库**:
导入torch、torchvision库用于处理深度学习相关的操作,以及数据预处理。
```python
import torch
import torchvision
from torchvision import datasets, transforms
```
2. **数据预处理**:
对MNIST数据进行归一化并加载到DataLoader中,方便批次训练。
```python
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
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 = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)
```
3. **构建模型**:
定义一个简单的CNN结构,例如包含几个卷积层、池化层和全连接层。
```python
class ConvNet(nn.Module):
def __init__(self):
super(ConvNet, self).__init__()
self.conv1 = nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv2 = nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1)
self.fc1 = nn.Linear(32 * 7 * 7, 128)
self.fc2 = nn.Linear(128, 10) # 输出层有10个节点对应MNIST的10个类别
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool(x)
x = F.relu(self.conv2(x))
x = self.pool(x)
x = x.view(-1, 32 * 7 * 7)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
model = ConvNet()
```
4. **设置损失函数和优化器**:
选择适合的损失函数(如交叉熵)和优化器(如SGD或Adam)。
```python
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
```
5. **训练循环**:
进行迭代训练,每次循环包括前向传播、计算损失、反向传播和更新权重。
```python
num_epochs = 10
for epoch in range(num_epochs):
for i, (images, labels) in enumerate(train_loader):
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad() # 清空梯度
loss.backward() # 反向传播
optimizer.step() # 更新权重
if (i+1) % 100 == 0:
print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{len(train_loader)}], Loss: {loss.item():.4f}")
```
6. **评估模型**:
在测试集上验证模型性能。
```python
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"Test Accuracy: {100 * correct / total}%")
```
阅读全文