本文采用3个卷积层、5个ReLu层、3个池化层和3个全连接层对MNIST数据集多分类。其中:
3个卷积层分别为(1, 16, kernel_size=5)、(16, 64, kernel_size=5)和(64, 256, kernel_size=3);3个池化层均为kernel_size=2×2的;3个全连接层分别为(256, 128)、(128, 64)、(64, 10)。代码如下:
# 0、导包
import torch
from torchvision import transforms,datasets
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.optim as optim
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
# 1、准备数据
batch_size = 64
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,),(0.3081,))
])
# 训练集
train_dataset = datasets.MNIST(root='G:/datasets/mnist',train=True,download=False,transform=transform)
train_loader = DataLoader(train_dataset,shuffle=True,batch_size=batch_size)
# 测试集
test_dataset = datasets.MNIST(root='G:/datasets/mnist',train=False,download=False,transform=transform)
test_loader = DataLoader(test_dataset,shuffle=False,batch_size=batch_size)
# 2.设计模型
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = torch.nn.Conv2d(1, 16, kernel_size=5)
self.conv2 = torch.nn.Conv2d(16, 64, kernel_size=5)
self.conv3 = torch.nn.Conv2d(64, 256, kernel_size=3)
self.pooling = torch.nn.MaxPool2d(2)
self.l1 = torch.nn.Linear(256, 128)
self.l2 = torch.nn.Linear(128, 64)
self.l3 = torch.nn.Linear(64, 10)
def forward(self, x):
# Flatten data from (n,1,28,28) to (n,784)
batch_size = x.size(0)
x = F.relu((self.conv1(x)))
x = self.pooling(x)
x = F.relu((self.conv2(x)))
x = self.pooling(x)
x = F.relu((self.conv3(x)))
x = self.pooling(x)
x = x.view(batch_size, -1) # flatten
x = F.relu(self.l1(x))
x = F.relu(self.l2(x))
return self.l3(x)
model = Net()
# 模型加载到GPU上
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
# 3、构造损失函数和优化器
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),lr=0.01,momentum=0.5)
# 4、训练和测试
def train(epoch):
running_loss = 0.0
for batch_idx, data in enumerate(train_loader, 0):
inputs, target = data
inputs, target = inputs.to(device), target.to(device)
optimizer.zero_grad()
# forward+backward+update
outputs = model(inputs)
loss = criterion(outputs, target)
loss.backward()
optimizer.step()
running_loss += loss.item()
if batch_idx % 300 == 299:
print('[%d,%d] loss: %.3f' % (epoch + 1, batch_idx + 1, running_loss / 300))
running_loss = 0.0
def test():
correct = 0
total = 0
with torch.no_grad():
for data in test_loader:
images, labels = data
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.data, dim=1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('Accuracy on test set: {0} %'.format(100 * correct / total))
return correct/total
accuracy_list = []
epoch_list = []
for epoch in range(10):
train(epoch)
accuracy_list.append(test())
epoch_list.append(epoch+1)
# 绘图
import matplotlib.pyplot as plt
plt.plot(epoch_list,accuracy_list, 'b')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.grid()
plt.show()
测试结果为:
[1,300] loss: 2.209
[1,600] loss: 0.625
[1,900] loss: 0.234
Accuracy on test set: 95.79 %
[2,300] loss: 0.150
[2,600] loss: 0.125
[2,900] loss: 0.099
Accuracy on test set: 97.7 %
[3,300] loss: 0.083
[3,600] loss: 0.079
[3,900] loss: 0.074
Accuracy on test set: 97.74 %
[4,300] loss: 0.060
[4,600] loss: 0.056
[4,900] loss: 0.062
Accuracy on test set: 98.55 %
[5,300] loss: 0.046
[5,600] loss: 0.050
[5,900] loss: 0.048
Accuracy on test set: 98.55 %
[6,300] loss: 0.038
[6,600] loss: 0.041
[6,900] loss: 0.039
Accuracy on test set: 98.84 %
[7,300] loss: 0.033
[7,600] loss: 0.036
[7,900] loss: 0.033
Accuracy on test set: 98.99 %
[8,300] loss: 0.029
[8,600] loss: 0.026
[8,900] loss: 0.031
Accuracy on test set: 98.87 %
[9,300] loss: 0.027
[9,600] loss: 0.025
[9,900] loss: 0.025
Accuracy on test set: 99.09 %
[10,300] loss: 0.021
[10,600] loss: 0.023
[10,900] loss: 0.023
Accuracy on test set: 99.06 %
[1] https://www.bilibili.com/video/BV1Y7411d7Ys?p=10