PyTorch實現(xiàn)手寫數(shù)字識別的示例代碼

更新時間：2022年05月31日 09:56:25 作者：B.Bz

本文主要介紹了PyTorch實現(xiàn)手寫數(shù)字識別的示例代碼，文中通過示例代碼介紹的非常詳細，具有一定的參考價值，感興趣的小伙伴們可以參考一下<BR>

加載手寫數(shù)字的數(shù)據(jù)

組成訓(xùn)練集和測試集，這里已經(jīng)下載好了，所以download為False

import torchvision

# 是否支持gpu運算
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# print(device)
# print(torch.cuda.is_available())


# 加載訓(xùn)練集的數(shù)據(jù)  使用torchvision自帶的MNIST數(shù)據(jù)集
train_dataset = torchvision.datasets.MNIST(root='./data1',
                                           train=True,
                                           transform=torchvision.transforms.ToTensor(),
                                           download=False
                                           )

# 加載測試集的數(shù)據(jù)  創(chuàng)建測試集
test_dataset = torchvision.datasets.MNIST(root='./data1',
                                          train=False,
                                          transform=torchvision.transforms.ToTensor(),
                                          download=False
                                          )

數(shù)據(jù)加載器（分批加載）

# 加載數(shù)據(jù)的批次 一批有多少條數(shù)據(jù)
batch_size = 100
# 創(chuàng)建數(shù)據(jù)加載器shuffle為True 加載時打亂
train_loader = DataLoader(dataset=train_dataset,
                          batch_size=batch_size,
                          shuffle=True
                          )
test_loader = DataLoader(dataset=test_dataset,
                         batch_size=batch_size,
                         shuffle=True
                         )
# 數(shù)據(jù)加載器生成的對象轉(zhuǎn)為迭代器
examples = iter(test_loader)
# 使用next方法獲取到一批次的數(shù)據(jù)
example_data, example_targets = examples.next()
# 遍歷獲取到6條數(shù)據(jù) 展示觀察一下
for i in range(6):
    plt.subplot(2, 3, i + 1)
    plt.imshow(example_data[i][0], cmap='gray')
    # 查看圖片的大小 方便建立模型時輸入的大小
    print(example_data[i][0].shape)

plt.show()

建立模型

建立模型之前定義輸入大小和分類類別輸出大小

通過上邊查看圖片的大小為28*28*1，所以輸入大小為784

數(shù)字識別只有0~9所以為10個類別的多分類問題

input_size = 784
num_classes = 10

創(chuàng)建模型類

class NeuralNet(torch.nn.Module):
    def __init__(self, n_input_size, hidden_size, n_num_classes):
        """
        神經(jīng)網(wǎng)絡(luò)類初始化
        :param n_input_size: 輸入
        :param hidden_size: 隱藏層
        :param n_num_classes: 輸出
        """
        # 調(diào)用父類__init__方法
        super(NeuralNet, self).__init__()
        self.input_size = input_size
        # 第一層線性模型 傳入輸入層和隱藏層
        self.l1 = torch.nn.Linear(n_input_size, hidden_size)
        # relu激活函數(shù)層
        self.relu = torch.nn.ReLU()
        # 第二層線性模型 傳入隱藏層和輸出層
        self.l2 = torch.nn.Linear(hidden_size, n_num_classes)

    def forward(self, x):
        """
        重寫正向傳播函數(shù)  獲取到預(yù)測值
        :param x: 數(shù)據(jù)
        :return: 預(yù)測值
        """
        # 線性模型
        out = self.l1(x)
        # 激活函數(shù)
        out = self.relu(out)
        # 線性模型2
        out = self.l2(out)
        # 返回預(yù)測值
        return out


# 獲取到gpu設(shè)備
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 創(chuàng)建模型并把模型放到當前支持的gpu設(shè)備中
model = NeuralNet(input_size, 500, num_classes).to(device)
print(model)

可以看出模型一共三層
輸入層（節(jié)點數(shù)量和圖小大小相同）
隱藏層（節(jié)點數(shù)為500）
輸出層（輸出節(jié)點數(shù)量為10 0~9）

定義損失函數(shù)和優(yōu)化器

因為是多分類問題，所以使用交叉熵函數(shù)的多分類損失函數(shù)
因為傳統(tǒng)的梯度下降存在一定缺陷，比如學(xué)習(xí)速率一直不變，所以使用PyTorch中梯度下降的優(yōu)化算法Adam算法

# 定義學(xué)習(xí)率
learning_rate = 0.01
# 損失函數(shù)
criterion = torch.nn.CrossEntropyLoss()
# 定義優(yōu)化器 參數(shù)1為模型的參數(shù) lr為學(xué)習(xí)率
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

模型訓(xùn)練

訓(xùn)練步驟:

通過模型類正向傳播獲取到預(yù)測結(jié)果
通過損失函數(shù)傳入預(yù)測結(jié)果和真實值計算損失
通過反向傳播獲取梯度
通過梯度下降更新模型參數(shù)的權(quán)重
梯度清空，防止下次梯度累加
循環(huán)，降低損失為我們想要的結(jié)果（提高模型精度）

# 定義訓(xùn)練的次數(shù)
num_epochs = 10
# 訓(xùn)練集數(shù)據(jù)的總長度
total_steps = len(train_loader)
# 遍歷訓(xùn)練次數(shù)
for epoch in range(num_epochs):
    # 每次從數(shù)據(jù)加載器中取出一批數(shù)據(jù)  每批次100條
    for i, (images, labels) in enumerate(train_loader):
        # 把圖片降維到一維數(shù)組  加載到gpu
        images = images.reshape(-1, 28 * 28).to(device)
        # 真實值加載到gpu
        labels = labels.to(device)
        # 正向傳播 獲取到預(yù)測值
        outputs = model(images)
        # 通過損失函數(shù)獲取到損失值
        loss_val = criterion(outputs, labels)
        # 清空梯度
        optimizer.zero_grad()
        # 進行反向傳播
        loss_val.backward()
        # 梯度下降更新參數(shù)
        optimizer.step()
        # 打印每次訓(xùn)練的損失值
        if i % 100 == 0:
            print(f'Loss:{loss_val.item():.4f}')

print('訓(xùn)練完成')
# 訓(xùn)練完之后保存模型
torch.save(model.state_dict(), './last.pt')

損失值很明顯的在收斂
生成了pt模型文件

測試集抽取數(shù)據(jù)，查看預(yù)測結(jié)果

# 把測試集的數(shù)據(jù)加載器轉(zhuǎn)為生成器
examples = iter(test_loader)
# next()方法獲取一批數(shù)據(jù)
example_data, example_targets = examples.next()

# 拿出前三條
for i in range(3):
    # 畫圖展示
    plt.subplot(1, 3, i + 1)
    plt.imshow(example_data[i][0], cmap='gray')
plt.show()

images = example_data
# 圖片將為加載到GPU
images = images.reshape(-1, 28 * 28).to(device)
# 正向傳播獲取預(yù)測結(jié)果
outputs = model(images)
# 打印結(jié)果 detach()方法結(jié)果不會計算梯度更新 轉(zhuǎn)為numpy
print(f'真實結(jié)果:{example_targets[0:3].detach().numpy()}')
# 預(yù)測完的結(jié)果為10個數(shù)字的概率 使用argmax()根據(jù)行歸一化并求自變量的概率最大值
print(f'預(yù)測結(jié)果:{np.argmax(outputs[0:3].cpu().detach().numpy(), axis=1)}')

計算模型精度

# 用測試集的數(shù)據(jù),校驗?zāi)Ｐ偷臏蚀_率
with torch.no_grad():
    n_correct = 0
    n_samples = 0
    # 取出測試集數(shù)據(jù)
    for images, labels in test_loader:
        # 和訓(xùn)練代碼一致
        images = images.reshape(-1, 28 * 28).to(device)
        labels = labels.to(device)
        outputs = model(images)

        # 返1 最大值 返2 索引                0每列最大值  1每行最大值
        _, predicted = torch.max(outputs.data, 1)
        n_samples += labels.size(0)
        n_correct += (predicted == labels).sum().item()
    # 計算模型精度
    acc = 100.0 * n_correct / n_samples
    print(f"準確率:{acc}%")

自己手寫數(shù)字進行預(yù)測

import cv2
import numpy as np

import torch

from 手寫數(shù)字神經(jīng)網(wǎng)絡(luò)結(jié)構(gòu) import NeuralNet

# 獲取到gpu設(shè)備
device = torch.device('cuda')
# 加載保存好的模型
input_size = 784
num_classes = 10
model = NeuralNet(input_size, 500, num_classes)
# 因為保存模型時在GPU所以要指定map_location='cuda:0'
model.load_state_dict(torch.load('./last.pt', map_location='cuda:0'))
# 加載到gpu上
model.to(device)

# 局域內(nèi)不計算梯度
with torch.no_grad():
    # cv2讀取圖片 灰度方式
    images = cv2.imread('./number_four.png', cv2.IMREAD_GRAYSCALE)
    # 使用大津算法進行二值化處理 并反轉(zhuǎn)
    ret, thresh_img = cv2.threshold(images, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
    # 展示處理過后的圖片
    cv2.imshow('png1', thresh_img)
    cv2.waitKey()
    # 圖片降維 把拍的圖片降維到和訓(xùn)練時的圖片大小一樣
    my_image = cv2.resize(thresh_img, (28, 28))
    # 轉(zhuǎn)為numpy
    my_image = np.array(my_image, np.float32)
    # 轉(zhuǎn)為torch的張量
    my_image = torch.from_numpy(my_image)
    # 降維
    my_image = my_image.reshape(-1, 28 * 28).to(device)
    # 正向傳播獲取預(yù)測值
    outputs = model(my_image)
    # 取出預(yù)測結(jié)果
    pred = np.argmax(outputs.cpu().detach().numpy(), axis=1)
    print(f'預(yù)測結(jié)果為：{pred[0]}')