专栏首页AI研习社如何用PyTorch训练图像分类器

如何用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)

原文出处及转载信息见文内详细说明,如有侵权,请联系 yunjia_community@tencent.com 删除。

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

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

我来说两句

0 条评论
登录 后参与评论

相关文章

  • 飞宇魏承赟:一个人的无人机江湖和他的九死一生

    「 唯物」按:有两句话在创业者中流传甚广,创业是一件九死一生的事情;选择比努力更重要。听上去就很悲壮和刺激。可现实往往不按套路出牌,看似同一个模子里刻出来的创业...

    AI研习社
  • 数据科学家必看!处理数据的 7 个小技巧

    在没有最基本的、必要的、未经处理数据的情况下,我们应该如何为机器学习的概念建立原型并加以验证呢?在资源匮乏的情况下,我们应如何有效地获取并用数据创造价值?

    AI研习社
  • 调试机器学习模型的六种方法

    在机器学习模型中,开发人员有时会遇到错误,但经常会在没有明确原因的情况下导致程序崩溃。虽然这些问题可以手动调试,但机器学习模型通常由于输出预测不佳而失败。更糟糕...

    AI研习社
  • 代码中两个常见的命名的概念

    Payload payload 在代码里可以理解为有效载重。 这个词来源于货运行业。比如有一位客户需要支付一笔费用委托货车司机运送一车石油,石油本身的重量、车子...

    linxinzhe
  • 一文梳理SPSS数据分析基础脉络,入门同学必备!

    “这是一场革命,庞大的数据资源使得各个领域开始了量化进程,无论是学术界、商界还是政府,所有领域都将开始这种进程。”

    博文视点Broadview
  • Centreon v19.04远程执行代码漏洞

    Centreon是一个免费的开源基础设施监控软件,Centreon允许系统管理员从集中式Web应用程序监控其基础设施,Centreon已成为欧洲企业监控...

    洛米唯熊
  • 如何测试OData的filter功能

    版权声明:本文为博主汪子熙原创文章,未经博主允许不得转载。 https://jerry.bl...

    Jerry Wang
  • python协程与golang协程的区

    进程,是计算机中已运行程序的实体。程序本身只是指令、数据及其组织形式的描述,进程才是程序的真正运行实例。

    py3study
  • 世界杯百亿盛宴暗潮涌动—到底是谁的足球狂欢?

    ? 世界杯盛宴——“狗庄”的秘密 足球比赛中,总有非法赌场庄家组织利用足球主客双方上下半场的进球、角球及边球的数量、比赛的胜负来进行外围赌博,他们起了个专业的...

    腾讯技术工程官方号
  • 【Web动画】CSS3 3D 行星运转 && 浏览器渲染原理

    Sb_Coco

扫码关注云+社区

领取腾讯云代金券