如何用PyTorch訓練圖像分類器

本文為 AI 研習社編譯的技術博客，原標題 ：

How to Train an Image Classifier in PyTorch and use it to Perform Basic Inference on Single Images

作者 | Chris Fotache

翻譯 | shunshun

校對 | 醬番梨 整理 | 菠蘿妹

原文鏈接：

https://medium.com/@chrisfotache/how-to-train-an-image-classifier-in-pytorch-and-use-it-to-perform-basic-inference-on-single-images-99465a1e9bf5

如果你剛剛開始使用PyTorch並想學習如何進行基本的圖像分類，那麼你可以參考本教程。它將介紹如何組織訓練數據，使用預訓練神經網路訓練模型，然後預測其他圖像。

為此，我將使用由Google地圖中的地圖圖塊組成的數據集，並根據它們包含的地形特徵對它們進行分類。我會在另一篇文章中介紹如何使用它（簡而言之：為了識別無人機起飛或降落的安全區域）。但是現在，我只想使用一些訓練數據來對這些地圖圖塊進行分類。

下面的代碼片段來自Jupyter Notebook。你可以將它們拼接在一起以構建自己的Python腳本，或從GitHub下載。這些Notebook是基於Udacity的PyTorch課程的。如果你使用雲端虛擬機進行深度學習開發並且不知道如何遠程打開notebook，請查看我的教程。

PyTorch希望數據按文件夾組織，每個類對應一個文件夾。大多數其他的PyTorch教程和示例都希望你先按照訓練集和驗證集來組織文件夾，然後在訓練集和驗證集中再按照類別進行組織。但我認為這非常麻煩，必須從每個類別中選擇一定數量的圖像並將它們從訓練集文件夾移動到驗證集文件夾。由於大多數人會通過選擇一組連續的文件作為驗證集，因此選擇可能存在很多偏差。

因此，這兒有一個將數據集快速分為訓練集和測試集的更好的方法，就像Python開發人員習慣使用sklearn一樣。首先，讓我們導入模塊：

%matplotlib inline

%config InlineBackend.figure_format = "retina"

import matplotlib.pyplot as plt

import numpy as np

import torch

from torch import nn

from torch import optim

import torch.nn.functional as F

from torchvision import datasets, transforms, models

接下來，我們將定義train/validation數據集載入器，使用SubsetRandomSampler進行拆分：

data_dir = "/data/train"

def load_split_train_test(datadir, valid_size = .2):

train_transforms = transforms.Compose([transforms.Resize(224),

transforms.ToTensor(),

])

test_transforms = transforms.Compose([transforms.Resize(224),

transforms.ToTensor(),

])

train_data = datasets.ImageFolder(datadir,

transform=train_transforms)

test_data = datasets.ImageFolder(datadir,

transform=test_transforms)

num_train = len(train_data)

indices = list(range(num_train))

split = int(np.floor(valid_size * num_train))

np.random.shuffle(indices)

from torch.utils.data.sampler import SubsetRandomSampler

train_idx, test_idx = indices[split:], indices[:split]

train_sampler = SubsetRandomSampler(train_idx)

test_sampler = SubsetRandomSampler(test_idx)

trainloader = torch.utils.data.DataLoader(train_data,

sampler=train_sampler, batch_size=64)

testloader = torch.utils.data.DataLoader(test_data,

sampler=test_sampler, batch_size=64)

return trainloader, testloader

trainloader, testloader = load_split_train_test(data_dir, .2)

print(trainloader.dataset.classes)

接下來我們將確定是否有GPU。我假設你有一台GPU機器，否則代碼將至少慢10倍。但是，檢查GPU可用性是個好主意。

我們還將載入預訓練模型。對於這種情況，我選擇ResNet 50：

device = torch.device("cuda" if torch.cuda.is_available

else "cpu")

model = models.resnet50(pretrained=True)

print(model)

列印模型將顯示ResNet模型的圖層體系結構。這可能超出了我的意識或你的理解，但看到那些深層隱藏層內的東西仍然很有趣。

這取決於你選擇什麼樣的模型，根據你的特定數據集模型可能會不同。這裡列出了所有的PyTorch模型。

現在我們進入深度神經網路的有趣部分。首先，我們必須凍結預訓練過的層，因此在訓練期間它們不會進行反向傳播。然後，我們重新定義最後的全連接層，即使用我們的圖像來訓練的圖層。我們還創建了標準（損失函數）並選擇了一個優化器（在這種情況下為Adam）和學習率。

for param in model.parameters:

param.requires_grad = False

model.fc = nn.Sequential(nn.Linear(2048, 512),

nn.ReLU,

nn.Dropout(0.2),

nn.Linear(512, 10),

nn.LogSoftmax(dim=1))

criterion = nn.NLLLoss

optimizer = optim.Adam(model.fc.parameters, lr=0.003)

model.to(device)

現在完成了，讓我們訓練模型吧！在這個例子中只有一個epoch，但在大多數情況下你需要更多。從代碼中可以看出基本過程非常直觀：載入批量圖像並執行前向傳播循環。然後計算損失函數，並使用優化器在反向傳播中應用梯度下降。

PyTorch就這麼簡單。下面的大多數代碼是每10個批次顯示損失並計算的準確度，所以你在訓練運行時得到更新。在驗證期間，不要忘記將模型設置為eval模式，然後在完成後返回train。

epochs = 1

steps = 0

running_loss = 0

print_every = 10

train_losses, test_losses = ,

for epoch in range(epochs):

for inputs, labels in trainloader:

steps += 1

inputs, labels = inputs.to(device), labels.to(device)

optimizer.zero_grad

logps = model.forward(inputs)

loss = criterion(logps, labels)

loss.backward

optimizer.step

running_loss += loss.item

if steps % print_every == 0:

test_loss = 0

accuracy = 0

model.eval

with torch.no_grad:

for inputs, labels in testloader:

inputs, labels = inputs.to(device),

labels.to(device)

logps = model.forward(inputs)

batch_loss = criterion(logps, labels)

test_loss += batch_loss.item

ps = torch.exp(logps)

top_p, top_class = ps.topk(1, dim=1)

equals =

top_class == labels.view(*top_class.shape)

accuracy +=

torch.mean(equals.type(torch.FloatTensor)).item

train_losses.append(running_loss/len(trainloader))

test_losses.append(test_loss/len(testloader))

print(f"Epoch {epoch+1}/{epochs}.. "

f"Train loss: {running_loss/print_every:.3f}.. "

f"Test loss: {test_loss/len(testloader):.3f}.. "

f"Test accuracy: {accuracy/len(testloader):.3f}")

running_loss = 0

model.train

torch.save(model, "aerialmodel.pth")

等待幾分鐘後（或更長時間後，取決於數據集的大小和時期數量），完成訓練並保存模型以供以後預測！

現在還有一件事可以做，即繪製訓練和驗證損失圖：

plt.plot(train_losses, label="Training loss")

plt.plot(test_losses, label="Validation loss")

plt.legend(frameon=False)

plt.show

如你所見，在我的一個epoch的特定例子中，驗證損失（這是我們感興趣的）在第一個epoch結束時的平坦線條甚至開始有上升趨勢，所以可能1個epoch就足夠了。正如預期的那樣，訓練損失非常低。

現在進入第二部分。你訓練模型，保存模型，並需要在應用程序中使用它。為此，你需要能夠對圖像執行簡單推理。你也可以在我們的存儲庫中找到此演示notebook。我們導入與訓練筆記本中相同的模塊，然後再次定義變換(transforms)。我只是再次聲明圖像文件夾，所以我可以使用那裡的一些例子：

data_dir = "/datadrive/FastAI/data/aerial_photos/train"

test_transforms = transforms.Compose([transforms.Resize(224),

transforms.ToTensor(),

])

然後我們再次檢查GPU可用性，載入模型並將其置於評估模式（因此參數不會改變）：

device = torch.device("cuda" if torch.cuda.is_available else "cpu")

model=torch.load("aerialmodel.pth")

model.eval

預測特定圖像的類的功能非常簡單。請注意，它需要Pillow圖像，而不是文件路徑。

def predict_image(image):

image_tensor = test_transforms(image).float

image_tensor = image_tensor.unsqueeze_(0)

input = Variable(image_tensor)

input = input.to(device)

output = model(input)

index = output.data.cpu.numpy.argmax

return index

現在為了便於測試，我還創建了一個從數據集文件夾中選擇大量隨機圖像的函數：

def get_random_images(num):

data = datasets.ImageFolder(data_dir, transform=test_transforms)

classes = data.classes

indices = list(range(len(data)))

np.random.shuffle(indices)

idx = indices[:num]

from torch.utils.data.sampler import SubsetRandomSampler

sampler = SubsetRandomSampler(idx)

loader = torch.utils.data.DataLoader(data,

sampler=sampler, batch_size=num)

dataiter = iter(loader)

images, labels = dataiter.next

return images, labels

最後，為了演示預測函數，我得到隨機圖像樣本，預測它們並顯示結果：

to_pil = transforms.ToPILImage

images, labels = get_random_images(5)

fig=plt.figure(figsize=(10,10))

for ii in range(len(images)):

image = to_pil(images[ii])

index = predict_image(image)

sub = fig.add_subplot(1, len(images), ii+1)

res = int(labels[ii]) == index

sub.set_title(str(classes[index]) + ":" + str(res))

plt.axis("off")

plt.imshow(image)

plt.show

以下是Google地圖圖塊上此類預測的一個示例。標籤是預測的類，我也在顯示它是否是正確的預測。

這就是它。繼續嘗試數據集。只要你正確組織圖像，此代碼應該按原樣運行。很快我就會有更多關於神經網路和PyTorch可以做的很酷的文章。

Chris Fotache是位於 New Jersey的 CYNET.ai的人工智慧研究員。他涵蓋了與生活中的人工智慧，Python編程，機器學習，計算機視覺，自然語言處理等相關的主題。雷鋒網雷鋒網雷鋒網

想要繼續查看該篇文章相關鏈接和參考文獻？

長按鏈接點擊打開或點擊【如何使用PyTorch訓練圖像分類器】：

http://ai.yanxishe.com/page/TextTranslation/1272

AI研習社每日更新精彩內容，觀看更多精彩內容：

使用Python來圖像增強

新手必看：手把手教你入門 Python

等你來譯：

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！