如何用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编程,机器学习,计算机视觉,自然语言处理等相关的主题。

原文发布于微信公众号 - AI研习社(okweiwu)

原文发表时间:2018-11-24

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏量化投资与机器学习

【年度系列】使用Tensorflow预测股票市场变动

我们将使用Tensorflow创建和开发一个简单的模型框架,以及提出一些对初步结果改进的意见。

20530
来自专栏大数据挖掘DT机器学习

R语言中的情感分析与机器学习

利用机器学习可以很方便的做情感分析。本篇文章将介绍在R语言中如何利用机器学习方法来做情感分析。在R语言中,由Timothy P.Jurka开发的情感分析以及更...

46130
来自专栏hadoop学习笔记

动态分配多任务资源的移动端深度学习框架

与云相比,移动系统受计算资源限制。然而众所周知,深度学习模型需要大量资源 。为使设备端深度学习成为可能,应用程序开发者常用的技术之一是压缩深度学习模型以降低其资...

12300
来自专栏机器之心

资源 | 可视化工具Yellowbrick:超参与行为的可视化带来更优秀的实现

17730
来自专栏AI研习社

微软开源 repo 1.0 ,旨在创造深度学习框架通用语言

AI 研习社按,日前,微软提出深度学习框架的通用语言——repo1.0,号称希望通过构建这一深度学习框架「Rosetta Stone(罗塞塔石碑)」,让研究者们...

14220
来自专栏人工智能

6步创建一个通用机器学习模板

小编说:本文将介绍一个通用的机器学习的项目模板,创建这个模板总共有六个步骤。你将会学到: •端到端地预测(分类与回归)模型的项目结构。 •如何将前面学到的内容引...

276100
来自专栏AI研习社

博客 | AI 从业者都应该知道的实验数据集

少了数据,我们的机器学习和深度学习模型什么也干不了。这么说吧,那些创建了数据集、让我们可以训练模型的人,都是我们的英雄,虽然这些人常常并没有得到足够的感谢。让人...

12920
来自专栏机器之心

百度NLP | 神经网络模型压缩技术

百度NLP专栏 作者:百度NLP 引言 近年来,我们在神经网络模型与 NLP 任务融合方面深耕,在句法分析、语义相似度计算、聊天生成等各类方向上,均取得显著的进...

46950
来自专栏数据科学学习手札

(数据科学学习手札13)K-medoids聚类算法原理简介&Python与R的实现

前几篇我们较为详细地介绍了K-means聚类法的实现方法和具体实战,这种方法虽然快速高效,是大规模数据聚类分析中首选的方法,但是它也有一些短板,比如在数据集中有...

51670
来自专栏QQ大数据团队的专栏

神盾推荐——离线算法平台

2.1K70

扫码关注云+社区

领取腾讯云代金券