如何用PyTorch训练图像分类器

标签: 人工智能 | 发表时间:2018-11-26 17:31 | 作者:
出处:https://www.leiphone.com

本文为 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

多目标追踪器:用OpenCV实现多目标追踪(C++/Python)

数据科学家应当了解的五个统计基本概念:统计特征、概率分布、降维、过采样/欠采样、贝叶斯统计


等你来译:

基于图像的路径规划:Dijkstra算法

掌握机器学习必须要了解的4个概念

正向和反向运动学:雅可比和微分运动

取得自然语言处理SOA结果的分层多任务学习模型(HMTL) 


相关 [pytorch 训练 图像] 推荐:

如何用PyTorch训练图像分类器

- - 雷锋网
本文为 AI 研习社编译的技术博客,原标题 :. 作者 |  Chris Fotache. 翻译 | shunshun. 校对 | 酱番梨        整理 | 菠萝妹. 如果你刚刚开始使用PyTorch并想学习如何进行基本的图像分类,那么你可以参考本教程. 它将介绍如何组织训练数据,使用预训练神经网络训练模型,然后预测其他图像.

字节跳动开源高性能分布式训练框架BytePS,支持PyTorch、TensorFlow等

- - IT瘾-tuicool
最近,字节跳动发布了一款通用高性能分布式训练框架 BytePS,该框架支持TensorFlow、Keras、PyTorch 和MXNet,并且可以在 TCP 或 RDMA 网络上运行. BytePS GitHub 地址:https://github.com/bytedance/byteps. 根据该项目的 GitHub 页面,BytePS 显著优于目前的开源分布式训练框架.

从Pipenv到PyTorch,盘点2017年最受欢迎的十大机器学习Python库

- -
参与:蒋思源、黄小天、刘晓坤. 2017 年即将结束,又到了总结的时刻. 本文作者把范围限定为机器学习,盘点了 2017 年以来最受欢迎的十大 Python 库;同时在这十个非常流行与强大的 Python 库之外,本文还给出了一些同样值得关注的 Python 库,如 PyVips 和 skorch.

使用Mxnet进行图像深度学习训练工具 InsightFace - 使用篇, 如何一键刷分LFW 99.80%, MegaFace 98%. - 知乎

- -
开头先把论文和开源项目地址放一下:. InsightFace库是对我们上述论文中提出算法的开源实现. 其中不仅仅包括我们自己的算法, 也包括其他常见的人脸loss, 比如Softmax, SphereFace, AMSoftmax, CosFace, Triplet Loss等等. 除了代码实现以外, 我们还提供了打包对齐好的人脸训练数据供下载, 免除了一大堆数据准备的工作.

Leetcode 编程训练

- - 酷 壳 - CoolShell.cn
Leetcode这个网站上的题都是一些经典的公司用来面试应聘者的面试题,很多人通过刷这些题来应聘一些喜欢面试算法的公司,比如:Google、微软、Facebook、Amazon之类的这些公司,基本上是应试教育的功利主义. 我做这些题目的不是为了要去应聘这些公司,而是为了锻炼一下自己的算法和编程能力.

我的听力训练之路

- moz - 学习力就是竞争力
舍得上初中那会,基本谈不上什么听力训练,最多最多就是英语老师拿了个录音机让大家做听写,听的通常是单词.就算如此,舍得初中四年的英语听力训练加在一起不会超过两个小时.. 初中毕业后舍得上的是技校,学的是热处理专业,根本就没有英语这一科,哈哈.技校毕业就进了工厂,玩了三年后才觉得该学学英语了.于是买了个爱华录音机,买了一套新概念英语的磁带,就开始听起来了.那时候年轻贪玩,加上基础差,没有坚持下来..

训练免疫系统杀死癌症

- 微笑!?~ - Solidot
一年前,当化疗对白血病也无效时,65岁的William Ludwig签署协议参与了宾夕法尼亚大学的一项大胆的实验. 医生抽取出10亿T细胞——对抗病毒和肿瘤的白细胞——为T细胞加入能编程细胞去攻击癌细胞的新基因,然后这些修改过的细胞被重新滴入Ludwig的静脉. 他浑身发冷颤,然后体温上升,血压下降,最后他被医生送入重病特别护理室,向其家人发出去世警告.

胜利还得靠“精神训练法”

- Clover - 果壳网 guokr.com - 果壳网
机会常常只有一次,只许成功不许失败. 你竭尽全力奔向目标,努力准备每一个细节,可是当机会来了,却发现准备有余,熟练不足. 明明已经写在笔记本上的要点,发言时一紧张又忘了说. 怎样才能把握住这些难得的表现机会,让自己熟练发挥呢. 对此,运动心理学提倡使用“精神训练法”(mental practice),简单说就是在脑子里把事情过一遍.

38种摄影基本功训练

- freefish - 佳人
摄影,除了拥有好装备,良好的运气,超凡的耐心之外,更有有一手过硬的摄影技巧,分享摄影基本功训练,文末提供摄影基本功训练全文阅读. 被摄体:一般风景、花卉、城市建筑等冲击力较强的景物. 建 议:首先使用广角镜头:24MM—35MM拍摄,光圈:F11—16,光圈优先AE模式. 要 求:只把焦点对在主要被摄体上,浅景深.

怎样训练写故事的能力?

- - 知乎每日精选
介是个好问题,又到了半夜写知乎的时间啦~. 我的以下答案,都是基于如何训练 戏剧性故事创作,用 戏剧创作的方法和经验来谈一谈故事创作,希望可以给题主一些帮助. 首先要说一说,戏剧性,啥叫戏剧性呢. 抛开资料论述,我个人最喜欢的一个解读为——. 那为什么要强调故事创作的戏剧性呢. 因为戏剧性满足了人自身对故事的渴望和期待习惯.