神经网络模型对手写数字的识别
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 plt
figure = 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 x
model = NeuralNetwork().to(device)
print(model)
def train(dataloader, model, loss_fn, optimizer):
model.train()
batch_size_num = 1
for 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, 0
with 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轮训练
- 测试结果如下: