Facebook技术分享：如何在PyTorch中训练图像分类模型

文章来源：企鹅号 - 猎维人工智能

图像分类是计算机视觉的最重要应用之一。它的应用范围包括从自动驾驶汽车中的物体分类到医疗行业中的血细胞识别，从制造业中的缺陷项目到建立可以对戴口罩的人进行分类的系统。在所有这些行业中，图像分类都以一种或另一种方式使用。他们是如何做到的呢？他们使用哪个框架？

你必须已阅读很多有关不同深度学习框架（包括TensorFlow，PyTorch，Keras等）之间差异的信息。TensorFlow和PyTorch无疑是业内最受欢迎的框架。下面我们将讲解如何在PyTorch中训练图像分类模型。

PyTorch概述

PyTorch在深度学习社区中越来越受欢迎，并且被深度学习从业者广泛使用，PyTorch是一个提供Tensor计算的Python软件包。此外，张量是多维数组，就像NumPy的ndarray一样，它们也可以在GPU上运行。

PyTorch的一个独特功能是它使用动态计算图。所述PyTorch的Autograd包从张量构建计算图和自动计算梯度。而不是具有特定功能的预定义图形。

PyTorch为我们提供了一个框架，可以随时随地构建计算图，甚至在运行时更改它们。特别是，对于我们不知道创建神经网络需要多少内存的情况，这很有用。你可以使用PyTorch应对各种深度学习挑战。以下是一些挑战：

图像（检测，分类等）

文字（分类，生成等）

强化学习

MNIST数据集

在开始之前，让我们了解数据集。在本文中，我们将解决流行的MNIST问题。这是一个数字识别任务，其中我们必须将手写数字的图像分类为0到9这10个类别之一。

在MNIST数据集中，我们具有从各种扫描的文档中获取的数字图像，其尺寸规格化并居中。随后，每个图像都是28 x 28像素的正方形（总计784像素）。数据集的标准拆分用于评估和比较模型，其中60,000张图像用于训练模型，而单独的10,000张图像集用于测试模型。

现在，我们也了解了数据集。因此，让我们在PyTorch和TensorFlow中使用CNN构建图像分类模型。我们将从在PyTorch中实现开始。我们将在google colab中实现这些模型，该模型提供免费的GPU以运行这些深度学习模型。

在PyTorch中实现CNN

让我们首先导入所有库：

我们还要在Google colab上检查PyTorch的版本：

因此，我正在使用1.5.1版本的PyTorch。如果您使用任何其他版本，则可能会收到一些警告或错误，因此您可以更新到此版本的PyTorch。我们将对图像执行一些转换，例如对像素值进行归一化，因此，让我们也定义这些转换：

现在，让我们加载MNIST数据集的训练和测试集：

接下来，我定义了训练和测试加载器，这将帮助我们分批加载训练和测试集。我将批量大小定义为64：

首先让我们看一下训练集的摘要：

因此，在每个批次中，我们有64个图像，每个图像的大小为28,28，并且对于每个图像，我们都有一个相应的标签。让我们可视化训练图像并查看其外观：

它是数字0的图像。类似地，让我们可视化测试集图像：

在测试集中，我们还有一批64号。现在让我们定义架构

定义模型架构

我们将在这里使用CNN模型。因此，让我们定义并训练该模型：

我们还定义优化器和损失函数，然后我们将看一下该模型的摘要：

因此，我们有2个卷积层，这将有助于从图像中提取特征。这些卷积层的要素传递到完全连接的层，该层将图像分类为各自的类别。现在我们的模型架构已经准备好，让我们训练这个模型十个时期：

可以看到训练随着时代的增加而减少。这意味着我们的模型是从训练集中学习的模式。让我们在测试集上检查该模型的性能：

因此，我们总共测试了10000张图片，并且该模型在预测测试图片的标签方面的准确率约为96％。这是ni 可以在PyTorch中构建卷积神经网络的方法.

扫码

添加站长进交流群

领取专属 10元无门槛券

私享最新 技术干货