莫凡pytorch mnist手写数据识别

mnist手写数据识别 CNNC训练

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt

#hyper parameters超参数
EPOCH = 1
BATCH_SIZE = 50
LR = 0.001
DOWNLOAD_MNIST = False

# 下载MNIST
train_data = torchvision.datasets.MNIST(
    root='./mnist',         #保存在文件夹中
    train=True,         #false则是测试点test
    transform=torchvision.transforms.ToTensor(),     #tensor数据的值在0-1之间 ，彩的0-255压缩到黑白的0-1
    download=DOWNLOAD_MNIST
    )
# #plot one example
# print(train_data.train_data.size())          #60000,28,28
# print(train_data.train_labels.size())      #60000
# plt.imshow(train_data.train_data[0].numpy(),cmap='gray')
# plt.title('%i'%train_data.train_labels[0])
# plt.show()                                           ctrl+/ duohangzhushi

train_loader = Data.DataLoader(dataset=train_data,batch_size=BATCH_SIZE,shuffle=True,num_workers=2)

test_data = torchvision.datasets.MNIST(root='./mnist/',train=False)
test_x = Variable(torch.unsqueeze(test_data.test_data,dim=1),volatile=True).type(torch.FloatTensor)[:2000]/255    #shape
test_y = test_data.test_labels[:2000]   #取前1000个标签y

#构建网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN,self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(                          #卷积层          (28,28,1)
                in_channels=1,               #图片的高度，黑白1层
                out_channels=16,          #输出的高度，过滤器的个数
                kernel_size=5,               #过滤器是5*5*16
                stride=1,                  #步长 1
                padding=2                  #填充 0   如果步长是1，则padding=(kernal_size-1)/2, 为了保证与前面图片大小相等
            ),                                                 # (28,28,16)
            nn.ReLU(),           #jihuohanshu
            nn.MaxPool2d(kernel_size=2),        #池化层    2*2  高度不变   (14,14,16)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16,32,5,1,2),               #加工成32层      (14,14,32)
            nn.ReLU(),
            nn.MaxPool2d(2),                                       #(7,7,32)
        )
        self.out = nn.Linear(32*7*7,10)        #将三维的数据展评成二维的数据

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)       #batch,327,7
        x = x.view(x.size(0),-1)      # batch,32*7*7
        output = self.out(x)          #输出10
        return output
cnn = CNN()
#print(cnn)

#训练网络
optimizer = torch.optim.Adam(cnn.parameters(),lr=LR)
loss_func = nn.CrossEntropyLoss()

for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):
        b_x = Variable(x)
        b_y = Variable(y)            # gives batch data, normalize x when iterate train_loader

        output = cnn(b_x)               # cnn output
        loss = loss_func(output, b_y)   # cross entropy loss
        optimizer.zero_grad()           # clear gradients for this training step
        loss.backward()                 # backpropagation, compute gradients
        optimizer.step()                # apply gradients


        if step % 50 ==0:
            test_output = cnn(test_x)
            pred_y = torch.max(test_output, 1)[1].data.squeeze()
            accuracy = sum(pred_y == test_y) / test_y.size(0)
            print('Epoch:',epoch,'| train loss: %.4f' % loss.item())

# print 10 predictions from test data
test_output = cnn(test_x[:10])
pred_y = torch.max(test_output, 1)[1].data.numpy().squeeze()
print(pred_y, 'prediction number')
print(test_y[:10].numpy(), 'real number')