手写数字识别从训练到部署全流程详解——Pytorch深度学习网络搭建与模型训练

综述：

自从alphaGo在围棋上战胜人类以来，以深度学习为核心的人工智能技术便得到了广泛的发展。其中算法设计与模型训练是深度学习研究的两个主要组成部分。在这里为了方便大家了解与入门深度学习开发，我以最简单的手写数字识别，通过Pytorch框架对LeNet-5网络进行构建与模型训练，，对Pytorch框架的使用与训练的全流程做一个详细的记录说明。
其中LeNet-5是一种用于手写体字符识别的非常高效的卷积神经网络，关于该网络的说明，大家可以通过论文（Gradient-Based Learning Applied to Document Recognition）和相关的技术博客（LeNet-5详解）来了解，这里不做过多赘述。

下面开始正题：

1. 训练框架环境

os: ubuntu 18.04
python version: 3.6+
torch 1.2+
torchvision 0.4+
opencv-python 
PIL
numpy 
onnx 
onnxruntime

2.网络搭建

# 搭建LeNet 网络模型
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 6, kernel_size=3, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(6, 16, kernel_size=5),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.fc1 = nn.Sequential(
            nn.Linear(16 * 5 * 5, 120),
            nn.BatchNorm1d(120),
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(120, 84),
            nn.BatchNorm1d(84),
            nn.ReLU()
        )
        self.fc3 = nn.Linear(84, 10)
        
    def forward(self, x):
        # print('x shape: ', x.shape)  # [N, 1, 28, 28]
        x = self.conv1(x)  # [N, 6, 14, 14]
        x = self.conv2(x)  # [N, 16, 5, 5]
        x = x.view(x.size()[0], -1)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

3.模型训练与导出

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    batch_size = 64
    # 下载和准备数据
    train_dataset = datasets.MNIST(root='./data/',
                                   train=True,
                                   transform=transforms.ToTensor(),
                                   download=True)
    test_dataset = datasets.MNIST(root='./data/',
                                  train=False,
                                  transform=transforms.ToTensor(),
                                  download=True)
    # 建立一个数据迭代器
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=batch_size,
                                               shuffle=True)
    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                              batch_size=batch_size,
                                              shuffle=False)

    net = LeNet().to(device)
    # 定义损失函数
    criterion = nn.CrossEntropyLoss()  
    # 设置训练参数
    LR = 0.001
    Momentum = 0.9
    optimizer = optim.SGD(net.parameters(), lr=LR, momentum=Momentum)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.8)
    epochs = 30
    
    epochs_acc = []
    for epoch in range(epochs):
        print("Epoch = ", epoch+1)
        # 训练模型
        sum_loss = 0.0
        for i, data in enumerate(train_loader):
            inputs, labels = data
            inputs, labels = Variable(inputs).cuda() if torch.cuda.is_available() else Variable(inputs).cpu(), \
                             Variable(labels).cuda() if torch.cuda.is_available() else Variable(labels).cpu()
            optimizer.zero_grad()#将梯度归零
            outputs = net(inputs)#将数据传入网络进行前向运算
            loss = criterion(outputs, labels)#得到损失函数
            loss.backward()#反向传播
            optimizer.step()#通过梯度做一步参数更新
            
            # print(loss)
            sum_loss += loss.item()
            if i % 100 == 99:
                print('[%d,%d] loss:%.03f' % (epoch + 1, i + 1, sum_loss / 100))
                sum_loss = 0.0
        scheduler.step()
                
        # 验证测试集
        net.eval()#将模型变换为测试模式
        correct = 0
        total = 0
        for data_test in test_loader:
            images, labels = data_test
            images, labels = Variable(images).cuda() if torch.cuda.is_available() else Variable(images).cpu(), \
                             Variable(labels).cuda() if torch.cuda.is_available() else Variable(labels).cpu()
            output_test = net(images)
            #此处的predicted获取的是最大值的下标
            _, predicted = torch.max(output_test, 1)
#             print("output_test:" + str(output_test))
            total += labels.size(0)
            correct += (predicted == labels).sum()
        print("correct sum: ", correct)
        epochs_acc.append(correct.item() / len(test_dataset))
        print("Test acc: {0}".format(correct.item() / len(test_dataset)))#.cpu().numpy()
        
        # 保存训练模型文件
        SAVE_PATH = "./Models/30/LeNet_p" + str(epoch+1) + ".pth"
        torch.save(net.state_dict(), SAVE_PATH)

    max_index, max_number = max(enumerate(epochs_acc), key=operator.itemgetter(1))
    print("Max acc epoch: ", max_index+1)

4.模型测试

4.1 pytorch模型测试

def pytorchModelTest(device, input, pytorch_model_path):
    # load pytorch model
    torch_model = LeNet().to(device)
    torch_model.load_state_dict(torch.load(pytorch_model_path, map_location=device))
    # set the model to inference mode
    torch_model.eval()
    # input to the model
    test_torch_out = torch_model(input)
    print("torch out:", test_torch_out)
    return test_torch_out

4.2 手写数字图片实测

4.2.1 测试代码

    # read img data
    img = Image.open(image_path).convert('L')
    test_transform = transforms.Compose([
        transforms.Resize(28),
        transforms.ToTensor()
    ])
    img2 = test_transform(img)
    img2 = torch.unsqueeze(img2, 0)
    img2 = img2.cuda().float() if torch.cuda.is_available() else img2.cpu().float()
    # pytorch test
    test_torch_out = pytorchModelTest(device, img2, input_pytorch_model_path)
    _, torch_predicted = torch.max(test_torch_out, 1)
    print("pytorch_test_img " + image_path + " out = ", torch_predicted)

4.2.2 实测效果图

输出结果：(输出tensor的序号代表识别到的数字)

torch out: tensor([[ 1.5921, 13.6434,  0.9059, -4.3882, -1.2459, -5.6476,  1.5759, -1.3486,
         -2.5218, -4.6718]], device='cuda:0', grad_fn=<AddmmBackward>)
pytorch_test_img ./data/test_hwd_imgs/1.jpg out =  tensor([1], device='cuda:0')

输出结果：(输出tensor的序号代表识别到的数字)

torch out: tensor([[-15.0987,  -6.8363,  -8.6192,   7.4078,   3.9232,  -3.8025,  -9.5131,
           2.4302,   5.4421,  26.3228]], device='cuda:0',
       grad_fn=<AddmmBackward>)
pytorch_test_img ./data/test_hwd_imgs/9.jpg out =  tensor([9], device='cuda:0')

完整代码

LeNet_Handwrite_Detect_Train_Demo