#AI挑战营第一站# MNIST手写数字识别

jianghelong

#AI挑战营第一站# MNIST手写数字识别 [复制链接]

本实验使用了由两层卷积层和一层全连接层组成的卷积神经网络实现了MNIST手写数字识别，得到的模型准确率为98.76%。

由于网络规模很小，本实验没有使用GPU，本实验使用的CPU为R5-7530u（六核十二线程，2GHz)

1、本实验构建的模型如下，完整代码见附件main.py：

class ConvolutionalNeuralWork(torch.nn.Module):
    def __init__(self, num_classes=10):
        super(ConvolutionalNeuralWork, self).__init__()
        self.conv_layer1 = torch.nn.Sequential(torch.nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1), torch.nn.ReLU(), torch.nn.MaxPool2d(kernel_size=2, stride=2))
        self.conv_layer2 = torch.nn.Sequential(torch.nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1), torch.nn.ReLU(), torch.nn.MaxPool2d(kernel_size=2, stride=2))
        self.linear_layer = torch.nn.Linear(1568, num_classes)

    def forward(self, x):
        x = self.conv_layer1(x)
        x = self.conv_layer2(x)
        x = x.reshape(x.size(0), -1)
        x = self.linear_layer(x)
        return x

2、训练代码如下，共训练了一百轮，完整代码见附件main.py：

wb = openpyxl.Workbook()
ws = wb.active
ws.cell(1, 1).value = "epoch"
ws.cell(1, 2).value = "train loss"
ws.cell(1, 3).value = "val accuracy"

num_epochs = 100
batch_size = 100
num_class = 10
learning_rate = 0.0001
#加载数据
train_dataset = torchvision.datasets.MNIST(root="./data0", train=True, transform=torchvision.transforms.ToTensor(), download=False)
test_dataset = torchvision.datasets.MNIST(root="./data0", train=False, transform=torchvision.transforms.ToTensor(), download=False)
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)

model = ConvolutionalNeuralWork(num_class)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

best_accuracy = 0
for epoch in range(num_epochs):
    loss_sum = 0
	#训练
    for image, label in train_loader:
        outputs = model(image)
        loss = criterion(outputs, label)
        loss_sum += loss.item()
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    print("epoch:", epoch)
    print("train loss:", loss_sum/train_dataset.data.shape[0])
    ws.cell(epoch+2, 1).value = epoch + 1
    ws.cell(epoch+2, 2).value = loss_sum/train_dataset.data.shape[0]
    total = 0
    correct = 0
	#测试
    for image, label in test_loader:
        outputs = model(image)
        s, predicted = torch.max(outputs, 1)
        total = label.size(0) + total
        for i in range(label.size(0)):
            if label[i] == predicted[i]:
                correct = correct + 1
    print("val accuracy:", correct / total)
    ws.cell(epoch+2, 3).value = correct / total
    if correct / total > best_accuracy:
        best_accuracy = correct / total
        print("better model")
        torch.save(model, "best.pth")
        ws.cell(epoch+2, 4).value = "better model"
    print('\n')
    wb.save("result.xlsx")

3、训练过程中训练集损失函数和测试集准确率如下所示，详细信息见附件result.xlsx

   

4、最终生成的模型文件*.pth，以及转换的ONNX模型见附件

秦天qintian0303

最终不是都需要弄到RV1106开发板上吗？