PyTorch入門-CIFAR10影象分類
阿新 • • 發佈：2020-08-11
CIFAR10資料集下載
CIFAR10資料集包含10個類別，影象尺寸為 3×32×32
官方下載地址很慢，這裡給一個百度雲 :
下載後在專案目錄新建一個data目錄解壓進去
匯入相關包
import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as transforms import time import copy MINI_BATCH = 8 # 資料集的圖片數量很大，無法一次性載入所有資料，所以一次載入一個mini-batch的圖片 DEVICE = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # GPU可用則使用GPU
使用torchvision載入並且歸一化訓練和測試資料集
CIFAR10資料集的輸出是範圍在[0,1]之間的PILImage，我們將它轉換並歸一化範圍在[-1,1]之間的Tensor
# ToTensor(): 將ndarrray格式的影象轉換為Tensor張量 # Normalize(std, mean): std：每個通道顏色平均值，這裡的平均值為0.5，私人資料集自己計算；mean：每個通道顏色標準偏差，(原始資料 - std) / mean 得到歸一化後的資料 transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
資料載入器
# 訓練資料載入 trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=False, transform=transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=MINI_BATCH, shuffle=True, num_workers=4) # 測試資料載入 testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=False, transform=transform) testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=4)
定義卷積神經網路
我們實現一個簡單的神經網路 LeNet-5來進行分類：
這個網路具有兩個卷積層，兩個池化層，三個全連線層，原網路用於手寫數字識別，輸入為灰度圖，這裡我們輸入影象是RGB所以修改輸入資料為 3×32×32 的Tensorr資料，輸出資料維度為 1*10 ，表示圖片屬於10個類別的概率，圖中資料維度變化說明：
- 二維卷積層輸出大小 out = (in - F + 2P) / S + 1 ，其中：
F： 卷積核大小 F×F
P： Padding，預設為0
S： 步長Stride，預設為1
如圖中第一層卷積層 (32 - 5) / 1 + 1 = 28 - 池化層輸出大小 ** out = (in - F) / S + 1** ，其中：
F： 池化視窗大小 F×F
S： 池化視窗移動的步長Stride，預設和池化視窗維度相同
如圖中第二層池化層 (28 - 2) / 2 + 1 = 14
這部分可以寫成一個獨立的檔案，在訓練程式碼中引入此檔案中的網路結構
# net.py
import torch
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5) # 卷積層：3通道到6通道，卷積5*5
self.conv2 = nn.Conv2d(6, 16, 5) # 卷積層：6通道到16通道，卷積5*5
self.pool = nn.MaxPool2d(2, 2) # 池化層，在2*2視窗上進行下采樣
# 三個全連線層 ：16*5*5 -> 120 -> 84 -> 10
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
# 定義資料流向
def forward(self, x):
x = F.relu(self.conv1(x)) # F.relu 是一個常用的啟用函式
x = self.pool(x)
x = F.relu(self.conv2(x))
x = self.pool(x)
x = x.view(-1, 16 * 5 * 5) # 變換資料維度為 1*(16*5*5)，-1表示根據後面推測
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
定義一個通用的訓練函式，每個epoch進行一次測試，最後返回具有最優引數的模型
def train(model, criterion, optimizer, epochs):
since = time.time()
best_acc = 0.0 # 記錄模型測試時的最高準確率
best_model_wts = copy.deepcopy(model.state_dict()) # 記錄模型測試出的最佳引數
for epoch in range(epochs):
print('-' * 30)
print('Epoch {}/{}'.format(epoch+1, epochs))
# 訓練模型
running_loss = 0.0
for i, data in enumerate(trainloader):
inputs, labels = data
inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
# 前向傳播，計算損失
outputs = net(inputs)
loss = criterion(outputs, labels)
# 反向傳播+優化
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss += loss.item()
# 每1000批圖片列印訓練資料
if (i != 0) and (i % 1000 == 0):
print('step: {:d}, loss: {:.3f}'.format(i, running_loss/1000))
running_loss = 0.0
# 每個epoch以測試資料的整體準確率為標準測試一下模型
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images.to(DEVICE), labels.to(DEVICE)
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
acc = correct / total
if acc > best_acc: # 當前準確率更高時更新
best_acc = acc
best_model_wts = copy.deepcopy(model.state_dict())
time_elapsed = time.time() - since
print('-' * 30)
print('訓練用時： {:.0f}m {:.0f}s'.format(time_elapsed//60, time_elapsed%60))
print('最高準確率: {}%'.format(100 * best_acc))
# 返回測試出的最佳模型
model.load_state_dict(best_model_wts)
return model
接下來只需要定義好損失函式和優化器然後呼叫訓練函式訓練模型即可
from net import Net
net = Net()
net.to(DEVICE)
# 使用分類交叉熵 Cross-Entropy 作損失函式，動量SGD做優化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# 訓練10個epoch
net = train(net, criterion, optimizer, 10)
# 儲存模型引數
torch.save(net.state_dict(), 'net_dict.pt')
測試模型
# 影象類別
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
net = Net()
net.load_state_dict(torch.load('net_dict.pt')) # 載入各層引數
net.to(DEVICE)
# 整體正確率
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images.to(DEVICE), labels.to(DEVICE)
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print('整體準確率: {}%'.format(100 * correct / total))
print('=' * 30)
# 每一個類別的正確率
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
with torch.no_grad():
for data in testloader:
images, labels = data
if torch.cuda.is_available():
images, labels = images.cuda(), labels.cuda()
outputs = net(images)
_, predicted = torch.max(outputs, 1)
c = (predicted == labels).squeeze()
for i in range(4):
label = labels[i]
class_correct[label] += c[i].item()
class_total[label] += 1
for i in range(10):
print('{}的準確率 : {:.2f}%'.format(classes[i], 100 * class_correct[i] / class_total[i]))
下面看一下模型對測試集圖片的一些預測結果
import matplotlib.pyplot as plt
import numpy as np
# 定義一個顯示圖片的函式
def imshow(img):
# 輸入資料：torch.tensor[c, h, w]
img = img / 2 + 0.5 # 反歸一
npimg = np.transpose(img.numpy(), (1, 2, 0)) # [c, h, w] -> [h, w, c]
plt.imshow(npimg)
plt.show()
# 取一批圖片
testdata = iter(testloader)
images, labels = testdata.next()
imshow(torchvision.utils.make_grid(images))
print('真實類別: ', ' '.join('{}'.format(classes[labels[j]]) for j in range(4)))
# 預測是10個標籤的權重，一個類別的權重越大，神經網路越認為它是這個類別，所以輸出最高權重的標籤。
outputs = net(images)
_, predicted = torch.max(outputs, 1)
print('預測結果: ', ' '.join('{}'.format(classes[predicted[j]]) for j in range(4)))