假设MINST数据集在E:\MNIST\train-images.idx3-ubyte中,怎么使用pytorch搭建一个七层CNN网络并使用这个MNIST数据集进行训练和测试,写出代码
时间: 2024-02-18 15:03:42 浏览: 87
好的,下面是使用PyTorch搭建一个七层CNN网络并使用MNIST数据集进行训练和测试的代码。注释中会对每个步骤进行说明:
```python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
# 定义超参数
batch_size = 64
learning_rate = 0.001
num_epochs = 10
# 定义数据预处理
transform = transforms.Compose([
transforms.ToTensor(), # 转换为Tensor类型
transforms.Normalize((0.1307,), (0.3081,)) # 标准化,使得均值为0,标准差为1
])
# 加载MNIST数据集
train_dataset = datasets.MNIST(root='E:/MNIST', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='E:/MNIST', train=False, transform=transform, download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
# 定义CNN模型
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
self.bn1 = nn.BatchNorm2d(32)
self.relu1 = nn.ReLU()
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
self.bn2 = nn.BatchNorm2d(64)
self.relu2 = nn.ReLU()
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
self.fc1 = nn.Linear(64 * 7 * 7, 128)
self.relu3 = nn.ReLU()
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
out = self.conv1(x)
out = self.bn1(out)
out = self.relu1(out)
out = self.conv2(out)
out = self.bn2(out)
out = self.relu2(out)
out = self.pool(out)
out = out.view(-1, 64 * 7 * 7)
out = self.fc1(out)
out = self.relu3(out)
out = self.fc2(out)
return out
# 实例化模型并定义损失函数和优化器
model = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# 训练模型
total_step = len(train_loader)
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()
# 每100个batch打印一次训练信息
if (i+1) % 100 == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
.format(epoch+1, num_epochs, i+1, total_step, loss.item()))
# 测试模型
model.eval() # 进入测试模式,关闭Dropout和BatchNormalization层
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('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))
```
这样,就可以使用PyTorch搭建一个七层CNN网络并使用MNIST数据集进行训练和测试了。
阅读全文