神经网络模型对手写数字的识别
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor """
MNIST包含70,000张手写数字图像:60,000张用于训练,10,000张用于测试。
图像是灰度的,28x28像素的,并且居中的,以减少预处理和加快运行。
"""
""" 下载训练数据集 (包含训练数据+标签)"""
training_data = datasets.MNIST(root='data',train=True,download=True,transform=ToTensor()
)
""" 下载测试数据集(包含训练图片+标签)"""
test_data = datasets.MNIST(root='data',train=False,download=True,transform=ToTensor()
)
print(len(training_data))""" 展示手写字图片 """
from matplotlib import pyplot as pltfigure = plt.figure()
for i in range(9):img, label = training_data[i + 59000] figure.add_subplot(3, 3, i + 1) plt.title(label)plt.axis("off") plt.imshow(img.squeeze(), cmap="gray")a = img.squeeze()
plt.show()training_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
for X, y in test_dataloader: print(f"Shape of X [N, C, H, W]: {X.shape}")print(f"Shape of y: {y.shape} {y.dtype}")break""" 判断当前设备是否支持GPU,其中mps是苹果m系列芯片的GPU """
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")class NeuralNetwork(nn.Module): def __init__(self):super().__init__() self.flatten = nn.Flatten() self.hidden1 = nn.Linear(28 * 28, 256) self.hidden2 = nn.Linear(256, 128) self.hidden3 = nn.Linear(128, 256)self.hidden4 = nn.Linear(256, 128)self.out = nn.Linear(128, 10)def forward(self, x): x = self.flatten(x)x = self.hidden1(x)x = torch.sigmoid(x) x = self.hidden2(x)x = torch.sigmoid(x)x = self.hidden3(x)x = torch.sigmoid(x)x = self.hidden4(x)x = torch.sigmoid(x)x = self.out(x)return xmodel = NeuralNetwork().to(device)
print(model)
def train(dataloader, model, loss_fn, optimizer):model.train() batch_size_num = 1for X, y in dataloader:X, y = X.to(device), y.to(device) pred = model.forward(X) loss = loss_fn(pred, y) optimizer.zero_grad() loss.backward() optimizer.step() loss_value = loss.item() if batch_size_num % 200 == 0:print(f"loss: {loss_value:>7f} [number:{batch_size_num}]")batch_size_num += 1
def test(dataloader, model, loss_fn):size = len(dataloader.dataset)num_batches = len(dataloader)model.eval() test_loss, correct = 0, 0with torch.no_grad(): for X, y in dataloader:X, y = X.to(device), y.to(device)pred = model.forward(X)test_loss += loss_fn(pred, y).item() correct += (pred.argmax(1) == y).type(torch.float).sum().item()a = (pred.argmax(1) == y) b = (pred.argmax(1) == y).type(torch.float)test_loss /= num_batches correct /= size print(f"Test result: \n Accuracy: {(100 * correct)}%, Avg loss: {test_loss}")loss_fn = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
epochs = 10
for e in range(epochs):print(f"Epoch {e + 1}\n")train(training_dataloader, model, loss_fn, optimizer)
print("Done!")
test(test_dataloader, model, loss_fn)
- 展示的手写数字图片如下:

- 模型结构如下:

- 训练结果如下:
- 共有10轮训练

- 测试结果如下:
