LeNet-1989复现

• Input：32 * 32 * 1的单值图像
• Convolution：kernel_size=5，Padding = 0， stride=0，out_channel=6，所以输出为28【32-5+1】 * 28 * 6
• Subsampling：池化，size=2，stride=2，所以输出为14*14【(28-2)/ 2+1】*6
• Convolution：kernel_size=5，Padding = 0， stride=0，out_channel=16，所以输出为10【14-5+1】 * 10 * 16
• Subsampling：池化，size=2，stride=2，所以输出为5*5【(10-2)/ 2+1】*16
• Fully Connected Layer：5516 -> 10 -> 84 ->10

import torch
import torchvision
import torch.nn as nn

from torch.utils.data import DataLoader
from torch.nn import functional as F
from torchvision import transforms
from matplotlib import pyplot as plt

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

n_epochs = 3
batch_size_train = 64
batch_size_test = 1000
learning_rate = 0.01
momentum = 0.5
log_interval = 10
random_seed = 1
torch.manual_seed(random_seed)
PATH = './LeNet.pth'


class LeNet(torch.nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)  # kernel_size = 5
        self.conv2 = nn.Conv2d(6, 16, 5)  # kernel_size = 5
        self.pool1 = nn.MaxPool2d(2, 2)  # kernel_size = 2, stride = 2

        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool1(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


def preprocess():
    transform = torchvision.transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)), transforms.Scale(32)])  # 缩放

    traindata = torchvision.datasets.MNIST('./data/', train=True, download=True, transform=transform)
    print('训练集数据：', traindata.data.shape)
    trainloader = torch.utils.data.DataLoader(traindata, batch_size=batch_size_train, shuffle=True)

    testdata = torchvision.datasets.MNIST('./data/', train=False, download=True, transform=transform)
    print('测试集数据：', testdata.data.shape)
    testloader = torch.utils.data.DataLoader(testdata, batch_size=batch_size_test, shuffle=True)

    return trainloader, testloader


def show(testloader):
    example = enumerate(testloader)
    id, (data, target) = next(example)
    print(data.shape)

    plt.imshow(data[0][0], cmap='gray')
    plt.show()


def train(trainloader):
    net = LeNet().cuda()

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=momentum)

    # 测试集数据
    example = enumerate(testloader)
    id, (test_x, test_y) = next(example)

    for epoch in range(4):  # 3个epoch
        for step, data in enumerate(trainloader):
            inputs, labels = data
            inputs, labels = inputs.cuda(), labels.cuda()

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            # print statistics
            if step % 50 == 0:
                test_output = net(test_x.cuda())
                pred_y = torch.max(test_output, 1)[1].cuda().data
                accuracy = torch.sum(pred_y.to(device) == test_y.to(device)).type(torch.FloatTensor) / test_y.size(0)
                print('Epoch: ', epoch, '| train loss: %.4f' % loss, '| test accuracy: %.2f' % accuracy)

    torch.save(net.state_dict(), PATH)  # save
    print('Finished Training')


def test(testloader):
    net = LeNet().cuda()
    net.load_state_dict(torch.load(PATH))

    all = 0
    acc = 0
    with torch.no_grad():  # 不track梯度
        for _, data in enumerate(trainloader):
            inputs, test_y = data
            inputs, test_y = inputs.cuda(), test_y.cuda()

            test_output = net(inputs.cuda())
            pred_y = torch.max(test_output, 1)[1].cuda().data
            acc += torch.sum(pred_y.to(device) == test_y.to(device)).type(torch.FloatTensor)
            all += test_y.size(0)

    print('在测试集上的Accuracy：', acc / all)


if __name__ == '__main__':
    trainloader, testloader = preprocess()
    # show(testloader)
    train(trainloader)
    test(testloader)