机器学习/深度学习--手写数字识别(MNIST数据集)

import torch
# 导入torchvision的transform模块，用来处理数据集
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

# 创建一个transforms实例，用于对数据集进行处理
transform = transforms.Compose([
    # 将图片转换为tensor
    transforms.ToTensor(),
    # 归一化
    transforms.Normalize((0.1307,), (0.3081,))
])

# 创建训练集
"""
datasets.MNIST是Pytorch的内置函数torchvision.datasets.MNIST，通过这个可以导入数据集。
    root设置数据集路径，如果该路径中有需要的数据，则使用，如果没有且download为True,则下载数据集到目录中。
    train=True 表明这是训练集
    transform=transform 表明使用上面定义的transform来处理数据集。
"""
# 创建训练集
train_dataset = datasets.MNIST(root='DataSet\\',
                               train=True,
                               download=True,
                               transform=transform)
# 创建测试集
test_dataset = datasets.MNIST(root='DataSet\\',
                              train=False,
                              download=True,
                              transform=transform)
# 定义训练数据装载器实例
train_loader = DataLoader(train_dataset,
                          batch_size=64,
                          shuffle=True)
# 定义测试数据装载器实例
test_loader = DataLoader(test_dataset,
                         batch_size=64,
                         shuffle=True)


# 定义模型
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # MNIST数据集中的图片是28*28=784的
        self.l1 = torch.nn.Linear(784, 512)
        self.l2 = torch.nn.Linear(512, 256)
        self.l3 = torch.nn.Linear(256, 128)
        self.l4 = torch.nn.Linear(128, 64)
        self.l5 = torch.nn.Linear(64, 10)

    def forward(self, x):
        # view中一个参数定为 - 1，代表自动调整这个维度上的元素个数，以保证元素的总数不变.每行784个元素
        # 如果设置为(-1,392)则x为2行392列
        x = x.view(-1, 784)
        x = torch.relu(self.l1(x))
        x = torch.relu(self.l2(x))
        x = torch.relu(self.l3(x))
        x = torch.relu(self.l4(x))
        return self.l5(x)

# 创建模型实例
net = Net()

# 设置损失函数和优化器
# 损失函数使用交叉熵
loss = torch.nn.CrossEntropyLoss(reduction='mean')
# 优化器使用随机梯度下降
optimizer = torch.optim.SGD(net.parameters(), lr=0.01, momentum=0.5)

# 定义训练
def train(epoch):
    running_loss = 0.0
    for data in train_loader:
        inputs, label = data
        y_pre = net(inputs)
        loss_num = loss(y_pre, label)
        # 累加每一次训练的损失
        running_loss += loss_num

        # 梯度置0
        optimizer.zero_grad()
        # 反向传播
        loss_num.backward()
        # 参数更新
        optimizer.step()

    # 输出一下将所有样本训练一遍的损失
    print("epoch:", epoch+1, "one-epoch-loss:", running_loss)


# 定义测试集
def test():
    correct = 0
    sum = 0
    # 在测试集中不需要进行梯度下架
    with torch.no_grad():
        for data in test_loader:
            inputs, label = data
            outputs = net(inputs)
            """
            torch.max(matrix,dim)
            dim有两个取值0和1,取0时表示按列索引matrix每列最大值，取1时表示按行索引matrix每行最大值
            返回值是两个torch,第一个values是找到的最大值，第二个indices是最大值的索引
            
            例如:
            a = torch.tensor([[1,5,62,54], [2,6,2,6], [2,65,2,6]])
            torch.max(a, 1)
            输出:
            torch.return_types.max(
                values=tensor([62, 6, 65]),
                indices=tensor([2, 3, 1]))
            
            """
            # "_",是占位符,接收torch.max()第一个返回值
            _, predicted = torch.max(outputs.data, dim=1)
            sum += label.size(0)
            correct += (predicted == label).sum().item()
        print("Accuracy on test:", correct / sum)


if __name__ == '__main__':
    for epoch in range(10):
        train(epoch)
        test()

经过十轮训练后，模型 识别准确率达到了97.8%

在测试集中将图片换成自己要预测的图片

def test():
    img = cv2.imread('F:\\pytorch_study\\test01\\DataSet\\num_pic\\7.png')
    img = cv2.resize(img, (28, 28))
    # 训练集是单通道的图片，这里将图片拆分，只取图片一个通道
    img1 = cv2.split(img)[0]
    img_tensor = transform(img1)
    data = img_tensor
    # 在测试集中不需要进行梯度下架
    with torch.no_grad():
        inputs = data
        outputs = net(inputs)
        # "_",是占位符,接收torch.max()第一个返回值
        _, predicted = torch.max(outputs.data, dim=1)
    print('预测结果：', predicted.data.item())

运行结果：