pytorch 多数字识别加定位（MNIST数据集）

学习过深度学习中的卷积神经网络的同志应该都接触过MNIST数据集，并且做过手写数字识别吧，还记得第一次做成功手写数字识别的那种成就感哈哈哈。

本深度学习小白最近尝试了对MNIST数据集进行了修改，将多个数字合成在一张图片中，用pyplot来显示，然后对其进行识别和定位。大概是这样的：

 比较简陋啊，但是大体上能满足要求就可以了。

具体过程

一  、导入所需要的库

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader,TensorDataset

二 、定义超参数

# 定义超参数
BATCH_SIZE = 16  # 每批次处理的数据
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 是否用GPU或者CPU
EPOCHS = 10  # 训练的数据集的轮次

三 、下载并加载数据

"""构建pipeline,对图像处理"""
pipeline = transforms.Compose([
    transforms.ToTensor(),  # 将图片转换成tensor
    transforms.Normalize((0.1307,), (0.3081,))  # 正则化 第一个参数是均值，第二个参数是标准差 这些系数都是数据集提供方计算好的数据
])

"""下载数据集 """

train_set = datasets.MNIST("data", train=True, download=True, transform=pipeline)

test_set = datasets.MNIST("data", train=False, download=True, transform=pipeline)


"""加载数据"""
train_loader = DataLoader(train_set, shuffle=True)

test_loader = DataLoader(test_set, shuffle=True)

四、制作训练集和测试集数据

""" 制作训练集数据和标签"""
number_x_train = np.zeros((30000, 1, 100, 100), dtype="float64")
coordinate_y_train = np.zeros((30000, 1, 4), dtype="float64")
number_y_train = np.zeros((30000, 1, 20), dtype="float64")

j = 0
for i, (data, label) in enumerate(train_loader):
    data = data.squeeze(axis=0)
    data = data.squeeze(axis=0)
    data = np.array(data)
    data[0, :], data[27, :], data[:, 0], data[:, 27] = 1, 1, 1, 1
    if i % 2 == 0:
        """背景板"""
        blank = np.zeros((100, 100))
        """临时存放标签和坐标 """
        multi_label = np.zeros((1, 20))
        multi_coordinate = np.array([[0, 0, 0, 0]])
        """只能生成在左半边的坐标"""
        dx = np.random.randint(0, 70)
        dy = np.random.randint(0, 20)
        multi_coordinate[0][0] = dx
        multi_coordinate[0][1] = dy
        """采用类似独热编码的方式存储 对应序号的数字为1 其余为0"""
        multi_label[0,int(label.numpy().item())] = 1
    else:
        """只能生成在右半边的坐标"""
        dx = np.random.randint(0, 70)
        dy = np.random.randint(50, 70)
        multi_coordinate[0][2] = dx
        multi_coordinate[0][3] = dy
        multi_label[0,int(label.numpy().item())+10] = 1
    """将数字放在背景板上"""
    blank[dx:dx + 28, dy:dy + 28] = blank[dx:dx + 28, dy:dy + 28] + data
    if i % 2 == 1:
        number_x_train[j, :, :, :] = blank
        number_y_train[j, :] = multi_label
        coordinate_y_train[j, :] = multi_coordinate
        plt.imshow(number_x_train[j].squeeze(0),cmap='binary')
        plt.show()
        j += 1

number_X_train = torch.from_numpy(number_x_train).float()
coordinate_Y_train = torch.from_numpy(coordinate_y_train).float()
number_Y_train = torch.from_numpy(number_y_train).float()


train_data = TensorDataset(number_X_train, number_Y_train)
train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)

就是每次都生成一个100*100 的背景板，对原MNIST数据集中的所有训练集和测试集中的所有数据，进行两两配对，并且随机生成坐标，将其显示在背景板上就可以了。我这里为了防止两个数字重叠，生成坐标的时候将其中一个固定在左半边，另一个固定在右半边。最后将这些数据保存下来就可以了。这里是训练集的代码，测试集几乎差不多也就不重复贴代码了。

五、识别数字模型

"""识别数字模型"""
class Multi_Digit(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=3,padding=1,stride=2)  # 1: 灰度图片的通道 10： 输出通道 3：kernel大小3*3
        self.conv2 = nn.Conv2d(10, 20, 3)  # 10: 输入通道 20： 输出通道 3：kernel大小3*3
        self.conv3 = nn.Conv2d(20, 40, 3)  # 10: 输入通道 20： 输出通道 3：kernel大小3*3
        self.fc1 = nn.Linear(3240, 500)  # 3240: 输入通道 500： 输出通道
        self.fc2 = nn.Linear(500, 20)  # 500:输入通道 20:输出通道

    def forward(self, x):  #  x 为  batch size *1 *100 *100   
        input_size = x.size(0)  # batch size
        x = self.conv1(x)  # 卷积   输入：batch *1 *100 *100 ,输出 ： batch*10*50*50
        x = F.relu(x)  # 激活函数 输出 ： batch*10*46*46
        x = F.max_pool2d(x, 2, 2)  # 池化层 f=2 s=2 输入 ：batch*10*50*50 输出 ：batch*10*25*25

        x = self.conv2(x)  # 卷积   输入：batch *10 *25 *25 ,输出 ： batch*20*23*23
        x = F.relu(x)
        x =F.max_pool2d(x,2,2)  # 卷积   输入：batch *20 *23 *23 ,输出 ： batch*20*11*11

        x =self.conv3(x) # 卷积   输入：batch *20 *11 *11 ,输出 ： batch*40*9*9
        x = F.relu(x)
        x = x.view(input_size, -1) 
        x = self.fc1(x)  # 输入 ：batch * 3240 输出：batch * 500
        x = F.relu(x)  
        x = self.fc2(x)  # 输入 batch *500   输出 ：batch *20
        return x

这个是识别数字的模型，因为图片也不是很复杂，所以这个模型也很简单。就卷积，激活，池化叠加起来就行了，最后全连接输出。

六、定义训练方法和测试方法

model = Multi_Digit().to(DEVICE)  # 创建模型部署到设备上 （CPU或者GPU）

optimizer = optim.Adam(model.parameters())  # 定义一个优化器 ，将模型的参数作为输入，优化参数

def train_model(model, device, train_loader, optimizer, epoch):  # train_loader为训练的数据
    # 数字识别模型训练
    model.train()  #启用batch normalization(批标准化)和drop out,model.train()是保证BN层能够用到每一批数据的均值和方差。对于Dropout，model.train()是随机取一部分网络连接来训练更新参数。
    for batch_index, (data, label) in enumerate(train_loader):  # data 为图片 label 为标签
        # 部署到DEVICE上去
        data, label = data.to(device), label.to(device)
        label = label.squeeze(axis=1)
        # 梯度初始化为 0
        optimizer.zero_grad() # 即将 optimizer中的weight置零
        # 训练后的结果
        output = model(data)
        # 计算损失
        MSE = nn.MSELoss()
        loss = MSE(output, label)
        # 反向传播
        loss.backward()
        # 参数优化
        optimizer.step()
        if batch_index % 10000 == 0:
            print("Train Epoch :{} \t Loss :{:.6f}".format(epoch, loss.item()))


# 定义数字识别测试方法
def test_model(model, device, test_loader):
    # 模型验证
    model.eval() 
    # 正确率
    correct = 0.0
    # 测试损失
    test_loss = 0.0
    n1,n2=0,0
    with torch.no_grad():  # 不会计算梯度，也不会反向传播
        for data, label in test_loader:
            # 部署到设备上
            data, label = data.to(device), label.to(device)
            # 测试数据
            output = model(data)
            # 计算测试损失
            label =label.squeeze(axis=1)
            MSE = nn.MSELoss()
            test_loss += MSE(output, label).item()
            # 累计正确率
            correct += np.round(output).eq(label.view_as(output)).sum().item()/20  
        test_loss /= len(test_loader.dataset)
        print("Test --- Average loss : {:.4f}  Accuracy:{:.4f}".format(test_loss,100.0 * correct / len(test_loader.dataset)))

#
for epoch in range(1, EPOCHS + 1):
    train_model(model, DEVICE, train_loader, optimizer, epoch)
    test_model(model, DEVICE, test_loader)

准确率还挺高的哈哈哈

七、定位模型

"""定位模型"""
class Multi_Number_Coordinate(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=3, padding=1, stride=2)  # 1: 灰度图片的通道 10： 输出通道 3：kernel大小3*3
        self.conv2 = nn.Conv2d(10, 20, 3)  # 10: 输入通道 20： 输出通道 3：kernel大小3*3
        self.conv3 = nn.Conv2d(20, 40, 3)  # 10: 输入通道 20： 输出通道 3：kernel大小3*3
        self.fc1 = nn.Linear(3240, 500)  # 8820: 输入通道 500： 输出通道
        self.fc2 = nn.Linear(500, 4)  # 500:输入通道 4:输出通道

    def forward(self, x):  # 前向传播  x 为  batch size *1 *100 *100  1是通道（灰色）
        input_size = x.size(0)  # batch size
        x = self.conv1(x)  # 卷积   输入：batch *1 *100 *100 ,输出 ： batch*10*50*50
        x = F.relu(x)  # 激活函数 输出 ： batch*10*46*46
        x = F.max_pool2d(x, 2, 2)  # 池化层 f=2 s=2 输入 ：batch*10*50*50 输出 ：batch*10*25*25

        x = self.conv2(x)  # 卷积   输入：batch *10 *25 *25 ,输出 ： batch*20*23*23
        x = F.relu(x)
        x = F.max_pool2d(x, 2, 2)  # 卷积   输入：batch *20 *23 *23 ,输出 ： batch*20*11*11

        x = self.conv3(x)  # 卷积   输入：batch *20 *11 *11 ,输出 ： batch*40*9*9
        x = F.relu(x)
        x = x.view(input_size, -1)  # 转化成一维 ， -1 激动计算维度

        x = self.fc1(x)  # 输入 ：batch * 8820 输出：batch * 500
        x = F.relu(x)  # 保持shape不变
        x = self.fc2(x)  # 输入 batch *500   输出 ：batch *10
        return x

定位模型几乎和识别数字模型差不多，一定有人问为什么还分开来两个模型，直接用一个不就好了嘛。的确可以这样，不过没试过这样准确率高不高。大家可以把修改的数据集将坐标和数字标签放在一起，然后只训练一个模型然后把最后全连接层改成24通道的就可以了。定位模型的训练方法和测试方法也和识别数字的差不多，这里也不重复贴代码了。

八、模型效果

然后就可以随便取出一组数据来输出看看效果啦。这准确率还算可以吧。

九、总结

虽然大概做出来了，但是要改进的应该也有很多，首先就是可以缩成一个模型。还有一些代码很笨重的地方。当然如果有其他修改意见欢迎提出来。之前一直纠结该用什么损失函数，那个时候不是用类似独热编码的方式保存标签的，只是单纯把数字记录下来，但这样似乎没有一个损失函数可以拿来用，后来才想到可以用这种方式，经验还是太少，以后接触的多了，应该会有更多想法吧。