Pytorch 手寫數(shù)字識別MNIST
內(nèi)容來自老師授課,為了方便個人學(xué)習(xí),自己手敲了一遍
老師的編寫的代碼有獨(dú)特的風(fēng)格,它將整個workflow分為幾三個文件,一個是模型文件(CNN.py),用來定義網(wǎng)絡(luò)結(jié)構(gòu)(模型),一個是訓(xùn)練集文件(train.py),用來訓(xùn)練模型,還有一個測試集文件(test.py),用來加載訓(xùn)練模型并在測試集上驗(yàn)證。





模型文件 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

訓(xùn)練模型 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
# 定義超參數(shù)
learning_rate = 1e-2 ?# 學(xué)習(xí)率
batch_size = 128 ?# 批處理的樣本大小
epoches_num = 20 ?# 遍歷訓(xùn)練集次數(shù)
# 下載訓(xùn)練集 MNIST 手寫數(shù)字識別
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()
# 開始訓(xùn)練
for epoch in range(epoches_num):
? ?print('*' * 40)
? ?running_loss = 0.0
? ?running_acc = 0.0
? ?# 訓(xùn)練
? ?for i, data in enumerate(train_loader, 1):
? ? ? ?img, label = data
? ? ? ?# GPU訓(xùn)練
? ? ? ?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')

測試并驗(yàn)證 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
# 定義超參數(shù)
batch_size = 128 ? ? # 批的大小
# 下載數(shù)據(jù)集 MNIST 手寫數(shù)字測試集
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}')

當(dāng)然,寫在同一個文件里也沒有任何問題,主要是熟悉deep learning中加載數(shù)據(jù)集、訓(xùn)練網(wǎng)絡(luò)、模型驗(yàn)證的工作流,作為初學(xué)者熟悉基本的操作。因?yàn)楝F(xiàn)有的深度學(xué)習(xí)框架已經(jīng)封裝了相應(yīng)的api,比如反向傳播、隨機(jī)梯度下降算法、損失函數(shù)、激活函數(shù)等等,可能不太care底層的算法,但是需要使用者自定義model,選擇超參數(shù),定義網(wǎng)絡(luò)結(jié)構(gòu)。這里還附上在知乎上學(xué)的一份代碼。

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) ?# 實(shí)驗(yàn)可重復(fù)
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()