Pytorch 手写数字识别MNIST
内容来自老师授课，为了方便个人学习，自己手敲了一遍
老师的编写的代码有独特的风格，它将整个workflow分为几三个文件，一个是模型文件（CNN.py），用来定义网络结构（模型），一个是训练集文件（train.py），用来训练模型，还有一个测试集文件（test.py）,用来加载训练模型并在测试集上验证。
模型文件 CNN.py
from torch import nn
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=3),
nn.BatchNorm2d(16),
nn.ReLU(inplace=True)
)
self.layer2 = nn.Sequential(
nn.Conv2d(16, 32, kernel_size=3),
nn.BatchNorm2d(32),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.layer3 = nn.Sequential(
nn.Conv2d(32, 64, kernel_size=3),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True)
)
self.layer4 = nn.Sequential(
nn.Conv2d(64, 128, kernel_size=3),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=2, stride=2)
)
self.fc = nn.Sequential(
nn.Linear(128 * 4 * 4, 1024),
nn.ReLU(inplace=True),
nn.Linear(1024, 128),
nn.ReLU(inplace=True),
nn.Linear(128, 10)
)
def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = x.reshape(x.size(0), -1)
fc_out = self.fc(x)
return fc_out
训练模型 train.py
import torch
import CNN
from torch import nn, optim
from torchvision import datasets
from torchvision import transforms
from torch.autograd import Variable
from torch.utils.data import DataLoader
# 定义超参数
learning_rate = 1e-2 # 学习率
batch_size = 128 # 批处理的样本大小
epoches_num = 20 # 遍历训练集次数
# 下载训练集 MNIST 手写数字识别
train_dataset = datasets.MNIST(root="./data", train=True, transform=transforms.ToTensor(), download=True)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# 定义model,loss,optimizer
model = CNN.SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
if torch.cuda.is_available():
print("CUDA is enable! GPU is available!")
model = model.cuda()
model.train()
# 开始训练
for epoch in range(epoches_num):
print('*' * 40)
running_loss = 0.0
running_acc = 0.0
# 训练
for i, data in enumerate(train_loader, 1):
img, label = data
# GPU训练
if torch.cuda.is_available():
img = torch.Tensor(img).cuda()
label = Variable(label).cuda()
else:
img = Variable(img)
label = Variable(label)
# 前向传播
out = model(img)
loss = criterion(out, label)
running_loss += loss.item() * label.size(0)
_, pred = torch.max(out, 1)
num_correct = (pred == label).sum()
accuracy = (pred == label).float().mean()
running_acc += num_correct.item()
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
print(
f"Finish {epoch + 1} Loss: {running_loss / len(train_dataset) :.6f}, Acc: {running_acc / len(train_dataset) :.6f}")
# 保存模型
torch.save(model, 'cnn.pt')
测试并验证 test.py
import torch
from torch import nn
from torchvision import datasets
from torchvision import transforms
from torch.autograd import Variable
from torch.utils.data import DataLoader
# 定义超参数
batch_size = 128 # 批的大小
# 下载数据集 MNIST 手写数字测试集
test_dataset = datasets.MNIST(root="./data", train=False, transform=transforms.ToTensor())
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
# 加载 Train 模型
model = torch.load("cnn.pt")
criterion = nn.CrossEntropyLoss()
model.eval()
eval_acc = 0
eval_loss = 0
# 测试
for data in test_loader:
img, label = data
if torch.cuda.is_available():
img = Variable(img).cuda()
label = Variable(label).cuda()
else:
img = Variable(img)
label = Variable(label)
out = model(img)
loss = criterion(out, label)
eval_loss += loss.item() * label.size(0)
_, pred = torch.max(out, 1)
num_correct = (pred == label).sum()
eval_acc += num_correct.item()
print(f'Test Loss:{eval_loss/len(test_dataset) :.6f}, Acc: {eval_acc/len(test_dataset) :.6f}')
当然，写在同一个文件里也没有任何问题，主要是熟悉deep learning中加载数据集、训练网络、模型验证的工作流，作为初学者熟悉基本的操作。因为现有的深度学习框架已经封装了相应的api，比如反向传播、随机梯度下降算法、损失函数、激活函数等等，可能不太care底层的算法，但是需要使用者自定义model，选择超参数，定义网络结构。这里还附上在知乎上学的一份代码。
ministpytorch.py
import torch
import torchvision
from torch.utils.data import DataLoader
# load dataset
n_epochs = 3
batch_size_train = 64
batch_size_test = 1000
learning_rate = 0.01
momentum = 0.5
log_interval = 10
random_seed = 1
torch.manual_seed(random_seed) # 实验可重复
train_loader = torch.utils.data.DataLoader(
torchvision.datasets.MNIST("./data/", train=True, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=batch_size_train, shuffle=True
)
test_loader = torch.utils.data.DataLoader(
torchvision.datasets.MNIST("./data/", train=False, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=batch_size_test, shuffle=True
)
examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)
# print(example_targets)
# print(example_data.shape)
import matplotlib.pyplot as plt
# fig = plt.figure()
# for i in range(6):
# plt.subplot(2, 3, i + 1)
# plt.tight_layout()
# plt.imshow(example_data[i][0], cmap="gray", interpolation='none')
# plt.title(f"Ground Truth: {example_targets[i]}")
# plt.xticks([])
# plt.yticks([])
# plt.show()
# create a neural network
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self,x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return F.log_softmax(x)
network = Net()
optimizer = optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum)
train_losses = []
train_counter = []
test_losses = []
test_counter = [i*len(train_loader.dataset) for i in range(n_epochs + 1)]
def train(epoch):
network.train()
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
output = network(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % log_interval == 0:
print(f"Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} "
f"{100. * batch_idx / len(train_loader) :.0f}%]\tLoss:{loss.item() :.6f}")
train_losses.append(loss.item())
train_counter.append(
(batch_idx*64) + ((epoch-1) * len(train_loader.dataset))
)
torch.save(network.state_dict(), './model.pth')
torch.save(optimizer.state_dict(), './optimizer.pth')
# train(1)
def test():
network.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
output = network(data)
test_loss += F.nll_loss(output, target, size_average=False).item()
pred = output.data.max(1, keepdim=True)[1]
correct += pred.eq(target.data.view_as(pred)).sum()
test_loss /= len(test_loader.dataset)
test_losses.append(test_loss)
print(f"\nTest set: Avg. loss:{test_loss :.4f}, Accuracy: {correct}/{len(test_loader.dataset)} "
f"({100. * correct / len(test_loader.dataset) :.0f}%)")
test()
for epoch in range(1,n_epochs + 1):
train(epoch)
test()
import matplotlib.pyplot as plt
fig = plt.figure
plt.plot(train_counter, train_losses, color="blue")
plt.scatter(test_counter, test_losses, color="red")
plt.legend(["Train_loss", "test_loss"], loc="upper right")
plt.xlabel('number of training examples seen')
plt.ylabel('negative log likehood loss')
plt.show()
examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)
with torch.no_grad():
output = network(example_data)
fig = plt.figure()
for i in range(6):
plt.subplot(2, 3, 1+i)
plt.tight_layout()
plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
plt.title(f'Prediction: {output.data.max(1, keepdim=True)[1][i].item()}')
plt.xticks([])
plt.yticks([])
plt.show()
continuted_network = Net()
continuted_optimizer = optim.SGD(network.parameters(), lr=learning_rate,momentum=momentum)
network_state_dict = torch.load('model.pth')
continuted_network.load_state_dict(network_state_dict)
optimizer_state_dict = torch.load('optimizer.pth')
continuted_optimizer.load_state_dict(optimizer_state_dict)
for i in range(4, 9):
test_counter.append(i*len(train_loader.dataset))
train(i)
test()
fig = plt.figure()
plt.plot(train_counter, train_losses, color="blue")
plt.scatter(test_counter, test_losses, color="red")
plt.legend(['Train Loss', 'Test Loss'], loc='upper right')
plt.xlabel('number of training examples seen')
plt.ylabel('negative log likelihood loss')
plt.show()
