0
| 本文作者： AI研習(xí)社-譯站 | 2018-11-26 17:31 |
本文為 AI 研習(xí)社編譯的技術(shù)博客，原標(biāo)題 ：
How to Train an Image Classifier in PyTorch and use it to Perform Basic Inference on Single Images
作者 | Chris Fotache
翻譯 | shunshun
校對(duì) | 醬番梨 整理 | 菠蘿妹
原文鏈接：
https://medium.com/@chrisfotache/how-to-train-an-image-classifier-in-pytorch-and-use-it-to-perform-basic-inference-on-single-images-99465a1e9bf5
如果你剛剛開始使用PyTorch并想學(xué)習(xí)如何進(jìn)行基本的圖像分類，那么你可以參考本教程。它將介紹如何組織訓(xùn)練數(shù)據(jù)，使用預(yù)訓(xùn)練神經(jīng)網(wǎng)絡(luò)訓(xùn)練模型，然后預(yù)測(cè)其他圖像。
為此，我將使用由Google地圖中的地圖圖塊組成的數(shù)據(jù)集，并根據(jù)它們包含的地形特征對(duì)它們進(jìn)行分類。我會(huì)在另一篇文章中介紹如何使用它（簡(jiǎn)而言之：為了識(shí)別無人機(jī)起飛或降落的安全區(qū)域）。但是現(xiàn)在，我只想使用一些訓(xùn)練數(shù)據(jù)來對(duì)這些地圖圖塊進(jìn)行分類。
下面的代碼片段來自Jupyter Notebook。你可以將它們拼接在一起以構(gòu)建自己的Python腳本，或從GitHub下載。這些Notebook是基于Udacity的PyTorch課程的。如果你使用云端虛擬機(jī)進(jìn)行深度學(xué)習(xí)開發(fā)并且不知道如何遠(yuǎn)程打開notebook，請(qǐng)查看我的教程。
PyTorch希望數(shù)據(jù)按文件夾組織，每個(gè)類對(duì)應(yīng)一個(gè)文件夾。大多數(shù)其他的PyTorch教程和示例都希望你先按照訓(xùn)練集和驗(yàn)證集來組織文件夾，然后在訓(xùn)練集和驗(yàn)證集中再按照類別進(jìn)行組織。但我認(rèn)為這非常麻煩，必須從每個(gè)類別中選擇一定數(shù)量的圖像并將它們從訓(xùn)練集文件夾移動(dòng)到驗(yàn)證集文件夾。由于大多數(shù)人會(huì)通過選擇一組連續(xù)的文件作為驗(yàn)證集，因此選擇可能存在很多偏差。
因此，這兒有一個(gè)將數(shù)據(jù)集快速分為訓(xùn)練集和測(cè)試集的更好的方法，就像Python開發(fā)人員習(xí)慣使用sklearn一樣。首先，讓我們導(dǎo)入模塊：
%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數(shù)據(jù)集加載器，使用SubsetRandomSampler進(jìn)行拆分：
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。我假設(shè)你有一臺(tái)GPU機(jī)器，否則代碼將至少慢10倍。但是，檢查GPU可用性是個(gè)好主意。
我們還將加載預(yù)訓(xùn)練模型。對(duì)于這種情況，我選擇ResNet 50：
device = torch.device("cuda" if torch.cuda.is_available()
else "cpu")
model = models.resnet50(pretrained=True)
print(model)
打印模型將顯示ResNet模型的圖層體系結(jié)構(gòu)。這可能超出了我的意識(shí)或你的理解，但看到那些深層隱藏層內(nèi)的東西仍然很有趣。
這取決于你選擇什么樣的模型，根據(jù)你的特定數(shù)據(jù)集模型可能會(huì)不同。這里列出了所有的PyTorch模型。
現(xiàn)在我們進(jìn)入深度神經(jīng)網(wǎng)絡(luò)的有趣部分。首先，我們必須凍結(jié)預(yù)訓(xùn)練過的層，因此在訓(xùn)練期間它們不會(huì)進(jìn)行反向傳播。然后，我們重新定義最后的全連接層，即使用我們的圖像來訓(xùn)練的圖層。我們還創(chuàng)建了標(biāo)準(zhǔn)（損失函數(shù)）并選擇了一個(gè)優(yōu)化器（在這種情況下為Adam）和學(xué)習(xí)率。
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)
現(xiàn)在完成了，讓我們訓(xùn)練模型吧！在這個(gè)例子中只有一個(gè)epoch，但在大多數(shù)情況下你需要更多。從代碼中可以看出基本過程非常直觀：加載批量圖像并執(zhí)行前向傳播循環(huán)。然后計(jì)算損失函數(shù)，并使用優(yōu)化器在反向傳播中應(yīng)用梯度下降。
PyTorch就這么簡(jiǎn)單。下面的大多數(shù)代碼是每10個(gè)批次顯示損失并計(jì)算的準(zhǔn)確度，所以你在訓(xùn)練運(yùn)行時(shí)得到更新。在驗(yàn)證期間，不要忘記將模型設(shè)置為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')
等待幾分鐘后（或更長(zhǎng)時(shí)間后，取決于數(shù)據(jù)集的大小和時(shí)期數(shù)量），完成訓(xùn)練并保存模型以供以后預(yù)測(cè)！
現(xiàn)在還有一件事可以做，即繪制訓(xùn)練和驗(yàn)證損失圖：
plt.plot(train_losses, label='Training loss')
plt.plot(test_losses, label='Validation loss')
plt.legend(frameon=False)
plt.show()
如你所見，在我的一個(gè)epoch的特定例子中，驗(yàn)證損失（這是我們感興趣的）在第一個(gè)epoch結(jié)束時(shí)的平坦線條甚至開始有上升趨勢(shì)，所以可能1個(gè)epoch就足夠了。正如預(yù)期的那樣，訓(xùn)練損失非常低。
現(xiàn)在進(jìn)入第二部分。你訓(xùn)練模型，保存模型，并需要在應(yīng)用程序中使用它。為此，你需要能夠?qū)D像執(zhí)行簡(jiǎn)單推理。你也可以在我們的存儲(chǔ)庫中找到此演示notebook。我們導(dǎo)入與訓(xùn)練筆記本中相同的模塊，然后再次定義變換(transforms)。我只是再次聲明圖像文件夾，所以我可以使用那里的一些例子：
data_dir = '/datadrive/FastAI/data/aerial_photos/train'
test_transforms = transforms.Compose([transforms.Resize(224),
transforms.ToTensor(),
])
然后我們?cè)俅螜z查GPU可用性，加載模型并將其置于評(píng)估模式（因此參數(shù)不會(huì)改變）：
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model=torch.load('aerialmodel.pth')
model.eval()
預(yù)測(cè)特定圖像的類的功能非常簡(jiǎn)單。請(qǐng)注意，它需要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
現(xiàn)在為了便于測(cè)試，我還創(chuàng)建了一個(gè)從數(shù)據(jù)集文件夾中選擇大量隨機(jī)圖像的函數(shù)：
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
最后，為了演示預(yù)測(cè)函數(shù)，我得到隨機(jī)圖像樣本，預(yù)測(cè)它們并顯示結(jié)果：
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地圖圖塊上此類預(yù)測(cè)的一個(gè)示例。標(biāo)簽是預(yù)測(cè)的類，我也在顯示它是否是正確的預(yù)測(cè)。
這就是它。繼續(xù)嘗試數(shù)據(jù)集。只要你正確組織圖像，此代碼應(yīng)該按原樣運(yùn)行。很快我就會(huì)有更多關(guān)于神經(jīng)網(wǎng)絡(luò)和PyTorch可以做的很酷的文章。
Chris Fotache是位于 New Jersey的 CYNET.ai的人工智能研究員。他涵蓋了與生活中的人工智能，Python編程，機(jī)器學(xué)習(xí)，計(jì)算機(jī)視覺，自然語言處理等相關(guān)的主題。雷鋒網(wǎng)雷鋒網(wǎng)雷鋒網(wǎng)
想要繼續(xù)查看該篇文章相關(guān)鏈接和參考文獻(xiàn)？
長(zhǎng)按鏈接點(diǎn)擊打開或點(diǎn)擊【如何使用PyTorch訓(xùn)練圖像分類器】：
http://ai.yanxishe.com/page/TextTranslation/1272
AI研習(xí)社每日更新精彩內(nèi)容，觀看更多精彩內(nèi)容：
使用Python來圖像增強(qiáng)
新手必看：手把手教你入門 Python
等你來譯：
雷峰網(wǎng)原創(chuàng)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。
