如何用PyTorch訓(xùn)練圖像分類器

本文作者： AI研習(xí)社-譯站

2018-11-26 17:31

導(dǎo)語：如果你剛剛開始使用PyTorch并想學(xué)習(xí)如何進行基本的圖像分類，那么你可以參考本教程。

本文為 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
校對 | 醬番梨整理 | 菠蘿妹
原文鏈接：
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í)如何進行基本的圖像分類，那么你可以參考本教程。它將介紹如何組織訓(xùn)練數(shù)據(jù)，使用預(yù)訓(xùn)練神經(jīng)網(wǎng)絡(luò)訓(xùn)練模型，然后預(yù)測其他圖像。

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

下面的代碼片段來自Jupyter Notebook。你可以將它們拼接在一起以構(gòu)建自己的Python腳本，或從GitHub下載。這些Notebook是基于Udacity的PyTorch課程的。如果你使用云端虛擬機進行深度學(xué)習(xí)開發(fā)并且不知道如何遠(yuǎn)程打開notebook，請查看我的教程。

組織訓(xùn)練數(shù)據(jù)集

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

因此，這兒有一個將數(shù)據(jù)集快速分為訓(xùn)練集和測試集的更好的方法，就像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進行拆分：

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è)你有一臺GPU機器，否則代碼將至少慢10倍。但是，檢查GPU可用性是個好主意。

我們還將加載預(yù)訓(xùn)練模型。對于這種情況，我選擇ResNet 50：

device = torch.device("cuda" if torch.cuda.is_available() 
                                  else "cpu")
model = models.resnet50(pretrained=True)
print(model)

打印模型將顯示ResNet模型的圖層體系結(jié)構(gòu)。這可能超出了我的意識或你的理解，但看到那些深層隱藏層內(nèi)的東西仍然很有趣。

這取決于你選擇什么樣的模型，根據(jù)你的特定數(shù)據(jù)集模型可能會不同。這里列出了所有的PyTorch模型。

現(xiàn)在我們進入深度神經(jīng)網(wǎng)絡(luò)的有趣部分。首先，我們必須凍結(jié)預(yù)訓(xùn)練過的層，因此在訓(xùn)練期間它們不會進行反向傳播。然后，我們重新定義最后的全連接層，即使用我們的圖像來訓(xùn)練的圖層。我們還創(chuàng)建了標(biāo)準(zhǔn)（損失函數(shù)）并選擇了一個優(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)練模型吧！在這個例子中只有一個epoch，但在大多數(shù)情況下你需要更多。從代碼中可以看出基本過程非常直觀：加載批量圖像并執(zhí)行前向傳播循環(huán)。然后計算損失函數(shù)，并使用優(yōu)化器在反向傳播中應(yīng)用梯度下降。

PyTorch就這么簡單。下面的大多數(shù)代碼是每10個批次顯示損失并計算的準(zhǔn)確度，所以你在訓(xù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')

等待幾分鐘后（或更長時間后，取決于數(shù)據(jù)集的大小和時期數(shù)量），完成訓(xùn)練并保存模型以供以后預(yù)測！

現(xiàn)在還有一件事可以做，即繪制訓(xùn)練和驗證損失圖：

plt.plot(train_losses, label='Training loss')
plt.plot(test_losses, label='Validation loss')
plt.legend(frameon=False)
plt.show()

如何用PyTorch訓(xùn)練圖像分類器

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

現(xiàn)在進入第二部分。你訓(xùn)練模型，保存模型，并需要在應(yīng)用程序中使用它。為此，你需要能夠?qū)D像執(zhí)行簡單推理。你也可以在我們的存儲庫中找到此演示notebook。我們導(dǎo)入與訓(xùn)練筆記本中相同的模塊，然后再次定義變換(transforms)。我只是再次聲明圖像文件夾，所以我可以使用那里的一些例子：

data_dir = '/datadrive/FastAI/data/aerial_photos/train'
test_transforms = transforms.Compose([transforms.Resize(224),
                                      transforms.ToTensor(),
                                     ])

然后我們再次檢查GPU可用性，加載模型并將其置于評估模式（因此參數(shù)不會改變）：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model=torch.load('aerialmodel.pth')
model.eval()

預(yù)測特定圖像的類的功能非常簡單。請注意，它需要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)在為了便于測試，我還創(chuàng)建了一個從數(shù)據(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ù)測函數(shù)，我得到隨機圖像樣本，預(yù)測它們并顯示結(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()