PyTorch卷积神经网络实现MNIST手写数字识别

直接run,数据集在执行train_dataset = datasets.MNIST(download = True)会自动下载，稍等就好

import torch
import torch.nn as nn
import torch.optim as optim

from torchvision import datasets, transforms

#define parameters
input_size = 28
num_class = 10
num_epochs = 3
batch_size = 64

#train data
train_dataset = datasets.MNIST(
    root = './data1',
    train = True,
    transform = transforms.ToTensor(),
    download = True
)

#test data
test_dataset = datasets.MNIST(
    root = './data1',
    train = False,
    transform = transforms.ToTensor()
)
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 = True
)

#build the net

class CNN(nn.Module):
    def __init__(self):
        super(CNN,self).__init__()
        self.conv1 = nn.Sequential(         #input size 1*28*28
            nn.Conv2d(
                in_channels = 1,            #grey
                out_channels = 16,          #output channel 16
                kernel_size = 5,            #the size of kernel 5*5
                stride = 1,
                padding = 2                 #output size 16*28*28
            ),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 2)   #Pooling, the output size is (16,14,14)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(
                in_channels = 16,
                out_channels = 32,
                kernel_size = 5,
                stride = 1,
                padding = 2
            ),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size = 2)  #output size 32*7*7
        )
        self.out = nn.Linear(32*7*7,10)    #torch.nn.Linear(input_size,hidden_size)

    def forward(self,x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0),-1)   #flatten：(batch_size, 32*7*7)
        output= self.out(x)
        return output

def accuracy(pre,labels):
    pred = torch.max(pre.data,1)[1]
    rights = pred.eq(labels.data.view_as(pred)).sum()

    return rights, len(labels)

#train net
net = CNN()

criterion = nn.CrossEntropyLoss()

optim = optim.Adam(net.parameters(), lr = 0.001)

for epoch in range(num_epochs):
    #save the result
    train_rights = []
    for batch_id, (data, target) in enumerate(train_loader):
        net.train()
        output = net(data)
        loss = criterion(output, target)
        optim.zero_grad()
        loss.backward()
        optim.step()
        right = accuracy(output,target)
        train_rights.append(right)
        if batch_id % 100 == 0:
            net.eval()
            val_rights = []

            for (data,target ) in test_loader:
                output = net(data)
                right = accuracy(output ,target)
                val_rights.append(right)
            train_r = (sum([tup[0] for tup in train_rights]),sum([tup[1] for tup in train_rights]))
            val_r = (sum([tup[0] for tup in val_rights]), sum([tup[1] for tup in val_rights]))
            #print(epoch)

            print('当前epoch:{}[{}/{} ){:.0f}%]\t损失:{:.6f}\t训练集准确率:{:.2f}%\t测试集准确率:{:.2f}%'.format(epoch,batch_id*batch_size,len(train_loader.dataset),
                                                                                                 100. * batch_id / len(train_loader), loss.data,
                                                                                                 100. * train_r[0] / train_r[1],
                                                                                                 100. *val_r[0]/val_r[1]))