PyTorch 实现 LeNet 分类（数据cifar10）

今天参考了-深度学习框架PyTorch:入门与实践，这本书自己也动手写写LeNet。

LeNet 网络结构如下：

PyTorch 使用 torchvision 自带的 CIFAR10 数据实现。

运行环境：pytorch 0.4.0 CPU版、Python 3.6、Windows 7

# Pytorch 0.4.0 LeNet实现cifar10分类
# @Time   :2018/6/14
# @Author :LiuYinxing

import torchvision as tv
import torchvision.transforms as transforms
import torch as t
import torch.nn as nn
import torch.nn.functional as F
from torch import optim
t.set_num_threads(8)


class Net(nn.Module):  # 网络结构
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        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 = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(x.size()[0], -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


def getData():  # 定义对数据的预处理
    transform = transforms.Compose([
        transforms.ToTensor(),  # 转为Tensor
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # 归一化
    ])

    trainset = tv.datasets.CIFAR10(root='/data/', train=True,  download=True, transform=transform)  # 训练集
    trainloader = t.utils.data.DataLoader(trainset, batch_size=4, shuffle=True)

    testset = tv.datasets.CIFAR10('/data/', train=False, download=True, transform=transform)  # 测试集
    testloader = t.utils.data.DataLoader(testset, batch_size=4, shuffle=False)
    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    return trainloader, testloader, classes


def trainModel():  # 训练模型
    trainloader, testloader, _ = getData()  # 获取数据
    net = Net()
    criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)  # 定义优化器

    for epoch in range(2):
        for step, (tx, ty) in enumerate(trainloader, 0):
            optimizer.zero_grad()  # 梯度清零
            py = net(tx)  # forward + backward
            loss = criterion(py, ty)
            loss.backward()
            optimizer.step()  # 更新参数
            if step % 1000 == 999: 
                acc = testNet(net, testloader)
                print('Epoch:', epoch, '|Step:', step, '|train loss:%.4f' % loss.item(), '|test accuracy:%.4f' % acc)

    print('Finished Training')
    return net


def testNet(net, testloader):  # 获取在测试集上的准确率
    correct, total = .0, .0
    for x, y in testloader:
        py = net(x)
        _, predicted = t.max(py, 1)  # 获取分类结果
        total += y.size(0)  # 记录总个数
        correct += (predicted == y).sum()  # 记录分类正确的个数
    return float(correct) / total


if __name__ == '__main__':
    net = trainModel()

注：

testloader = t.utils.data.DataLoader(
                    testset,
                    batch_size=4, 
                    shuffle=False,
                    num_workers=2)

num_workers=2  # 多线程来读数据，在Windows使用时应该特别注意，很容易出错。

欢迎指正哦。