Tsubasa 的日记

卷积神经网络是近年发展起来， 并引起广泛重视的一种高效识别方法。

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

EPOCH = 3 # 训练三次
BATCH_SIZE = 50 # 一次训练的样本数量
LR = 0.001 # 学习率

train_data = torchvision.datasets.MNIST(
            root = '../data',
            train = True,
            transform = torchvision.transforms.ToTensor(),
            download = False
        )
test_data = torchvision.datasets.MNIST(
            root = '../data',
            train = False,
            transform = torchvision.transforms.ToTensor()
        )

train_loader = Data.DataLoader(
            dataset = train_data,
            batch_size = BATCH_SIZE,
            shuffle = True
        )
test_loader = Data.DataLoader(
            dataset = train_data,
            batch_size = BATCH_SIZE,
            shuffle = False
        )

class CNN(nn.Module) :
    def __init__(self) :
        super(CNN, self).__init__()
        self.layer1 = nn.Sequential( 
                    nn.Conv2d(
                            in_channels = 1,
                            out_channels = 16,
                            kernel_size = 5,
                            stride = 1, #　步长
                            padding = 2, # 补0
                            groups = 1
                        ),
                    nn.ReLU(),
                    nn.MaxPool2d(kernel_size = 2)
                )
        self.layer2 = nn.Sequential(
                    nn.Conv2d(
                            in_channels = 16,
                            out_channels = 32,
                            kernel_size = 5,
                            stride = 1, # 步长
                            padding = 2 # 补0
                        ),
                    nn.ReLU(),
                    nn.MaxPool2d(kernel_size = 2)
                )
        self.layer3 = nn.Linear(32 * 7 * 7, 10)
    def forward(self, x) :
        x = self.layer1(x)
        x = self.layer2(x)
        x = x.view(x.size(0), -1)
        x = self.layer3(x)
        return x

cnn = CNN().cuda()

optimizer = torch.optim.Adam(cnn.parameters(), lr = LR)
loss = nn.CrossEntropyLoss()

for epoch in range(EPOCH) :
    for step, (b_x, b_y) in enumerate(train_loader) :
        output = cnn(Variable(b_x).cuda())
        loss_ = loss(output, Variable(b_y).cuda())
        optimizer.zero_grad()
        loss_.backward()
        optimizer.step()

print("训练成功")

cnn.eval()
correct = 0
total = 0
for images, labels in test_loader :
    image = Variable(images).cuda()
    labels = Variable(labels).cuda()
    output = cnn(image).cuda()
    _, predicted = torch.max(output.data, 1)
    total += labels.size(0)
    correct += (predicted.cpu() == labels.cpu()).sum()

print("测试准确率 : %f %%" % (100 * correct / total))

最后输出结果为

训练成功
测试准确率 : 99.000000 %

比逻辑回归的正确率高出 18%!