torch系列之手写数字识别（实战）

公众号：道长的log

识别准确率99%，

有问题留言。 

直接上代码：

文件夹名称：MNIST_CNN_train

文件：model_CNN.py

from torch import nn
import torch.nn.functional as F


# model
class MnistModel(nn.Module):
    def __init__(self):
        super(MnistModel, self).__init__()

        # Convolution layer 1
        #  channels 卷积核的数量
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=5, stride=1, padding=2)
        self.relu1 = nn.ReLU()
        self.batch1 = nn.BatchNorm2d(64)

        self.conv2 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=5, stride=1, padding=2)
        self.relu2 = nn.ReLU()
        self.batch2 = nn.BatchNorm2d(128)
        # 最大池化层
        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.drop1 = nn.Dropout(0.25)

        # Convolution layer 2
        self.conv3 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1)
        self.relu3 = nn.ReLU()
        self.batch3 = nn.BatchNorm2d(256)

        self.conv4 = nn.Conv2d(in_channels=256, out_channels=128, kernel_size=3, stride=1, padding=1)
        self.relu4 = nn.ReLU()
        self.batch4 = nn.BatchNorm2d(128)
        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.drop2 = nn.Dropout(0.25)

        self.conv5 = nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.relu5 = nn.ReLU()
        self.batch5 = nn.BatchNorm2d(64)
        self.drop3 = nn.Dropout(0.25)

        # Fully-Connected layer 1

        self.fc1 = nn.Linear(3136, 256)
        self.fc1_relu = nn.ReLU()
        self.batch5 = nn.BatchNorm2d(64)
        self.dp1 = nn.Dropout(0.25)

        # Fully-Connected layer 2
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        # conv layer 1 的前向计算，3行代码
        out = self.conv1(x)
        out = self.relu1(out)
        out = self.batch1(out)

        out = self.conv2(out)
        out = self.relu2(out)
        out = self.batch2(out)

        out = self.maxpool1(out)
        out = self.drop1(out)

        # conv layer 2 的前向计算，4行代码
        out = self.conv3(out)
        out = self.relu3(out)
        out = self.batch3(out)

        out = self.conv4(out)
        out = self.relu4(out)
        out = self.batch4(out)

        out = self.maxpool2(out)
        out = self.drop2(out)

        out = self.conv5(out)
        out = self.relu5(out)
        out = self.batch5(out)
        out = self.drop3(out)

        # Flatten拉平操作
        out = out.view(out.size(0), -1)

        # FC layer的前向计算（2行代码）
        out = self.fc1(out)
        out = self.fc1_relu(out)
        out = self.dp1(out)

        out = self.fc2(out)

        return F.log_softmax(out, dim=1)

文件：train.py

from torch import save, load
from torchvision.datasets import MNIST
from torchvision import transforms
from torch.utils.data import DataLoader
from torch import nn
from torch import optim
from tqdm import tqdm
from MNIST_CNN_train import test
from MNIST_CNN_train import model_CNN
import numpy as np
import os
from torchsummary import summary

# 实例化模型
model = model_CNN.MnistModel()
print(model)
summary(model, (1, 28, 28))
optimizer = optim.Adam(model.parameters())

# 加载已经训练好的模型和优化器继续进行训练
if os.path.exists('./models/model.pkl'):
    model.load_state_dict(load("./models/model.pkl"))
    optimizer.load_state_dict(load("./models/optimizer.pkl"))

loss_function = nn.CrossEntropyLoss()
my_transforms = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.1307, ), std=(0.3081, ))
     ]
)
mnist_train = MNIST(root="../MNIST_data", train=True, download=True, transform=my_transforms)


def train(epoch):
    total_loss = []
    dataloader = DataLoader(mnist_train, batch_size=32, shuffle=True)
    dataloader = tqdm(dataloader, total=len(dataloader))
    model.train()
    for images, labels in dataloader:
        # 梯度置0
        optimizer.zero_grad()
        # 前向传播
        output = model(images)
        # 通过结果计算损失
        loss = loss_function(output, labels)
        total_loss.append(loss.item())
        # 反向传播
        loss.backward()
        # 优化器更新
        optimizer.step()

    save(model.state_dict(), './models/model.pkl')
    save(optimizer.state_dict(), './models/optimizer.pkl')
    # 打印一下训练成功率
    print('第{}个epoch，成功率, 损失为{}'.format(epoch+1, np.mean(total_loss)), test.test_success())


for i in range(66):
    train(i)

文件：test.py

from torchvision.datasets import MNIST
from torchvision import transforms
from torch.utils.data import DataLoader
from torch import nn
from tqdm import tqdm
import torch
import numpy as np
from MNIST_CNN_train import model_CNN
import os


def test_success():
    total_loss = []
    # 实例化模型
    model = model_CNN.MnistModel()
    if os.path.exists("./models/model.pkl"):
        model.load_state_dict(torch.load("./models/model.pkl"))
    loss_function = nn.CrossEntropyLoss()
    my_transforms = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(mean=(0.1307, ), std=(0.3081, ))
        ]
    )
    mnist_train = MNIST(root="../MNIST_data", train=False, download=True, transform=my_transforms)
    dataloader = DataLoader(mnist_train, batch_size=8, shuffle=True)
    dataloader = tqdm(dataloader, total=len(dataloader))

    succeed = []

    model.eval()
    with torch.no_grad():
        for images, labels in dataloader:
            # 获取结果
            output = model(images)
            result = output.max(dim=1).indices
            print("labels", labels)
            print("result", result)
            succeed.append(result.eq(labels).float().mean().item())
            # 通过结果计算损失
            loss = loss_function(output, labels)
            total_loss.append(loss.item())
    print(np.mean(total_loss))
    return np.mean(succeed)


if __name__ == '__main__':
    print(test_success())