一个简单的pytorch神经网络训练手写数字识别数据集

只有七十行代码，可以很方便学习者理解神经网络的整个训练过程。

代码：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from tqdm import tqdm

# Define the neural network
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28*28, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 28*28)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# Load the data
train_data = MNIST(root='./data', train=True, transform=ToTensor(), download=True)
test_data = MNIST(root='./data', train=False, transform=ToTensor(), download=True)
train_loader = DataLoader(train_data, batch_size=64, shuffle=True)
test_loader = DataLoader(test_data, batch_size=64, shuffle=False)

# Define the device
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# Define the model and optimizer
model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=0.01)

# Define the loss function
criterion = nn.CrossEntropyLoss()

# Train the model
for epoch in range(10):
    running_loss = 0.0
    correct = 0
    total = 0
    for i, data in tqdm(enumerate(train_loader), total=len(train_loader)):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    print(f'Epoch {epoch+1} loss: {running_loss/len(train_loader):.4f} accuracy: {100*correct/total:.2f}%')

# Test the model
model.eval()
test_loss = 0.0
correct = 0
total = 0
with torch.no_grad():
    for data in tqdm(test_loader, total=len(test_loader)):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        test_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
print(f'Test loss: {test_loss/len(test_loader):.4f} accuracy: {100*correct/total:.2f}%')

训练结果

Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./data/MNIST/raw/train-images-idx3-ubyte.gz
100%|█████████████████████████████████████████████████████| 9912422/9912422 [00:04<00:00, 2319644.71it/s]Extracting ./data/MNIST/raw/train-images-idx3-ubyte.gz to ./data/MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./data/MNIST/raw/train-labels-idx1-ubyte.gz
100%|██████████████████████████████████████████████████████████| 28881/28881 [00:00<00:00, 909782.30it/s]Extracting ./data/MNIST/raw/train-labels-idx1-ubyte.gz to ./data/MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw/t10k-images-idx3-ubyte.gz
100%|██████████████████████████████████████████████████████| 1648877/1648877 [00:02<00:00, 715696.00it/s]Extracting ./data/MNIST/raw/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw

Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz
100%|██████████████████████████████████████████████████████████| 4542/4542 [00:00<00:00, 21190799.52it/s]Extracting ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw

100%|█████████████████████████████████████████████████████████████████| 938/938 [00:07<00:00, 133.33it/s]Epoch 1 loss: 1.2479 accuracy: 74.10%
100%|█████████████████████████████████████████████████████████████████| 938/938 [00:05<00:00, 157.03it/s]Epoch 2 loss: 0.4967 accuracy: 87.40%
100%|█████████████████████████████████████████████████████████████████| 938/938 [00:05<00:00, 156.87it/s]Epoch 3 loss: 0.3944 accuracy: 89.27%
100%|█████████████████████████████████████████████████████████████████| 938/938 [00:05<00:00, 157.40it/s]Epoch 4 loss: 0.3529 accuracy: 90.09%
100%|█████████████████████████████████████████████████████████████████| 938/938 [00:06<00:00, 155.22it/s]Epoch 5 loss: 0.3284 accuracy: 90.68%
100%|█████████████████████████████████████████████████████████████████| 938/938 [00:05<00:00, 159.92it/s]Epoch 6 loss: 0.3103 accuracy: 91.22%
100%|█████████████████████████████████████████████████████████████████| 938/938 [00:05<00:00, 158.22it/s]Epoch 7 loss: 0.2957 accuracy: 91.61%
100%|█████████████████████████████████████████████████████████████████| 938/938 [00:05<00:00, 158.63it/s]Epoch 8 loss: 0.2830 accuracy: 92.00%
100%|█████████████████████████████████████████████████████████████████| 938/938 [00:05<00:00, 158.75it/s]Epoch 9 loss: 0.2715 accuracy: 92.34%
100%|█████████████████████████████████████████████████████████████████| 938/938 [00:05<00:00, 156.80it/s]Epoch 10 loss: 0.2615 accuracy: 92.63%
100%|█████████████████████████████████████████████████████████████████| 157/157 [00:00<00:00, 171.43it/s]Test loss: 0.2482 accuracy: 92.99%