Pytorch之VGG16图像分类

  • 个人主页:风间琉璃
  • 版权: 本文由【风间琉璃】原创、在CSDN首发、需要转载请联系博主
  • 如果文章对你有帮助、欢迎关注、点赞、收藏(一键三连)订阅专栏

目录

一、VGG

1.VGG网络结构

(1)输入层

(2)第一层卷积层

(3)第二层卷积层

(4)第三层卷积层

(5)第四层卷积层

(6)第五层卷积层

(7)第一层全连接层

(8)第二层全连接层

(9)第三层全连接层

(10)softmax层

2.全卷积网络

(1)第一层全连接层

(2)第二层全连接层

(3)第三层全连接层

(4)全卷积层作用

3.VGG创新点

1.网络层数深

2.小卷积核

3.全卷积层

二、VGG16实现

1.定义VGG网络模型

2.加载数据集

3.训练模型

4.测试模型

三、实现图像分类


一、VGG

VGG是Oxford的Visual Geometry Group的组提出的。该网络是在ILSVRC 2014上的相关工作,主要工作是证明了增加网络的深度能够在一定程度上影响网络最终的性能。斩获该年ImageNet竞赛中Localization Task(定位任务)第一名和Classification Task(分类任务)第二名。

1.VGG网络结构

VGG 的结构与 AlexNet 类似,区别是深度更深,但形式上更加简单。VGG由5层卷积层、3层全连接层、1层softmax输出层构成,层与层之间使用maxpool(最大化池)分开,所有隐藏层的激活单元都采用ReLU函数。作者在原论文中,根据卷积层不同的子层数量,设计了A、A-LRN、B、C、D、E这6种网络结构,如下图所示,

 Pytorch之VGG16图像分类_第1张图片

这6种网络结构相似,都是由5层卷积层、3层全连接层组成,区别在于每个卷积层的子层数量不同,从A至E依次增加,总的网络深度从11层到19层。

表格中的卷积层参数表示为“conv(卷积核大小)-通道数”,例如conv3-64,表示使用3x3的卷积核,通道数为64;最大池化表示为maxpool,层与层之间使用maxpool分开;全连接层表示为“FC-神经元个数”,例如FC-4096表示包含4096个神经元的全连接层;最后是softmax层。

虽然作者给了6个VGG网络的不同配置,尝试了不同的深度(11、13、16、19层)以及是否采用LRN等,但是在实际使用过程中,一般都会采用D这个配置,即16层:13个卷积层以及最后3个全连接层。

VGG16网络结构如下所示:

Pytorch之VGG16图像分类_第2张图片

其处理过程:

注意:

1.conv的stride为1,padding为1通过3*3的卷积核,输入、输出尺寸不变

2.maxpooling的size为2,stride为2通过maxpool,将特征矩阵的高和宽直接缩小一半

(1)输入层

输入224x224x3大小的RGB图像

(2)第一层卷积层

第1层卷积层由2个conv3-64和maxpooling组成,该层的处理流程是:卷积-->ReLU--> 卷积-->ReLU-->池化

卷积:输入是224x224x3,使用64个3x3x3的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size - kernel_size +  2 * padding) / stride + 1=(224+2*1-3)/1+1=224

得到输出是224x224x64

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

卷积输入是224x224x64,使用64个3x3x64的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(224+2*1-3)/1+1=224

得到输出是224x224x64。

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

池化使用2x2,stride=2的池化单元进行最大池化操作(max pooling)。根据公式:

(224+2*0-2)/2+1=112

最终得到的输出为112x112x64

(3)第二层卷积层

第2层卷积层由2个conv3-128和maxpooling组成。该层的处理流程是:卷积-->ReLU--> 卷积-->ReLU-->池化

卷积输入是112x112x64,使用128个3x3x64的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(112+2*1-3)/1+1=112

得到输出是112x112x128。

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

卷积输入是112x112x128,使用128个3x3x128的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(112+2*1-3)/1+1=112

得到输出是112x112x128。

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

池化:使用2x2,stride=2的池化单元进行最大池化操作(max pooling)。根据公式:

(112+2*0-2)/2+1=56

最终得到的输出为56x56x128

(4)第三层卷积层

第3层卷积层由3个conv3-256组成。该层的处理流程是:卷积-->ReLU--> 卷积-->ReLU-->卷积-->ReLU-->池化

卷积输入是56x56x128,使用256个3x3x128的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(56+2*1-3)/1+1=56

得到输出是56x56x256。

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

卷积输入是56x56x256,使用256个3x3x256的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(56+2*1-3)/1+1=56

得到输出是56x56x256。

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

卷积输入是56x56x256,使用256个3x3x256的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(56+2*1-3)/1+1=56

得到输出是56x56x256。

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

池化使用2x2,stride=2的池化单元进行最大池化操作(max pooling)。根据公式:

(56+2*0-2)/2+1=28

最终得到的输出为28x28x256

(5)第四层卷积层

第4层卷积层由3个conv3-512组成。该层的处理流程是:卷积-->ReLU--> 卷积-->ReLU--> 卷积-->ReLU-->池化

卷积输入是28x28x256,使用512个3x3x256的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(28+2*1-3)/1+1=28

得到输出是28x28x512。

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

卷积输入是28x28x512,使用512个3x3x512的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(28+2*1-3)/1+1=28

得到输出是28x28x512。

卷积输入是28x28x512,使用512个3x3x512的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(28+2*1-3)/1+1=28

得到输出是28x28x512。

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

池化使用2x2,stride=2的池化单元进行最大池化操作(max pooling)。根据公式:

(28+2*0-2)/2+1=14

最终得到的输出为14x14x512

(6)第五层卷积层

第5层卷积层由3个conv3-512组成。该层的处理流程是:卷积-->ReLU--> 卷积-->ReLU--> 卷积-->ReLU-->池化

卷积输入是14x14x512,使用512个3x3x512的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(14+2*1-3)/1+1=14

得到输出是14x14x512。

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

卷积输入是14x14x512,使用512个3x3x512的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(14+2*1-3)/1+1=14

得到输出是14x14x512。

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

卷积输入是14x14x512,使用512个3x3x512的卷积核进行卷积,padding=1,stride=1,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(14+2*1-3)/1+1=14

得到输出是14x14x512。

ReLU将卷积层输出的FeatureMap输入到ReLU函数中。

池化使用2x2,stride=2的池化单元进行最大池化操作(max pooling)。根据公式:

(14+2*0-2)/2+1=7

最终得到的输出为7x7x512

(7)第一层全连接层

第1层全连接层FC4096由4096个神经元组成。该层的处理流程是:FC-->ReLU-->Dropout

FC:输入是7x7x512的FeatureMap,展开为7*7*512的一维向量,即7*7*512个神经元,输出为4096个神经元。

ReLU:这4096个神经元的运算结果通过ReLU激活函数中。

Dropout:随机的断开全连接层某些神经元的连接,通过不激活某些神经元的方式防止过拟合。

(8)第二层全连接层

第2层全连接层FC4096由4096个神经元组成。该层的处理流程是:FC-->ReLU-->Dropout

FC:输入是4096个神经元,输出为4096个神经元。

ReLU:这4096个神经元的运算结果通过ReLU激活函数中。

Dropout:随机的断开全连接层某些神经元的连接,通过不激活某些神经元的方式防止过拟合。

(9)第三层全连接层

第3层全连接层FC1000由1000个神经元组成,对应ImageNet数据集的1000个类别

该层的处理流程是:FC

FC:输入是4096个神经元,输出为1000个神经元

(10)softmax层

该层的流程为:Softmax

Softmax:这1000个神经元的运算结果通过Softmax函数中,输出1000个类别对应的预测概率值

各层卷积核如下:

Pytorch之VGG16图像分类_第3张图片

2.全卷积网络

在VGG16后三层全连接网络中,VGG16在训练的时候使用的是全连接网络。然而在测试验证阶段,网络结构稍有不同,作者将全连接全部替换为卷积网络

(1)第一层全连接层

输入为7x7x512的FeatureMap,使用4096个7x7x512的卷积核进行卷积,由于卷积核尺寸与输入的尺寸完全相同,即卷积核中的每个系数只与输入尺寸的一个像素值相乘一一对应,根据公式:

(input_size + 2 * padding - kernel_size) / stride + 1=(7+2*0-7)/1+1=1

得到输出是1x1x4096,相当于4096个神经元,但属于卷积层

(2)第二层全连接层

输入为1x1x4096的FeatureMap,使用4096个1x1x4096的卷积核进行卷积,得到输出是1x1x4096,属于卷积层。

(3)第三层全连接层

输入为1x1x4096的FeatureMap,使用1000个1x1x4096的卷积核进行卷积,得到输出是1x1x1000。得到1x1x1000的输出之后,最后经过softmax层进行预测类别。

(4)全卷积层作用

将三个全连接层转成了1个7×7和 2 个 1×1 的卷积层。从下图可以看到,以第一个全连接层为例,要转卷积层,FC6的输入是 7×7×512,输出是4096(1×1×4096),那么就要对输入在尺寸上(宽高)降维(从7×7 降到 1×1)和深度(channel 或者 depth)升维(从512 升到4096)。

把7×7降到1×1,使用大小为 7×7的卷积核,卷积核个数设置为4096,即卷积核为7×7×4096(下图中的[7×7×512]×4096 表示有 4096 个 [7×7×512] 这样的卷积核,经过对输入卷积就得到了最终的 1×1×4096 大小的 feature map。

Pytorch之VGG16图像分类_第4张图片

全连接层转变为全卷积层原因让网络模型可以接受任意大小的尺寸,大大减少特征位置对分类带来的影响

网络输入图片的尺寸是224x224x3。如果后面三个层都是全连接,遇到宽高大于224的图片就需要进行图片的剪裁、缩放或其它处理,使图片尺寸统一到224x224x3,才能符合后面全连接层的输入要求。但是,我们并不能保证每次裁剪都能将图片中的关键目标保留下来,可能裁剪去的部分恰好包含了目标,造成裁减丢失关键目标信息,影响模型的测试精度。

输出是一个分类得分图,通道的数量和类别的数量相同,空间分辨率依赖于输入图像尺寸。最终为了得到固定尺寸的分类得分向量,将分类得分图进行空间平均化(求和——池化)。我们同样使用水平翻转来对测试集进行数据增强;在原始图像和翻转图像上的soft-max分类概率的平均值作为这幅图像的最终得分。

使用全卷积层,即使图片尺寸大于224x224x3,最终经过softmax层得到的得分图就不是1x1x1000,例如是2x2x1000,这里的通道1000与类别的数量相同,空间分辨率2x2依赖于输入图像尺寸。然后将得分图进行空间平均化(求和池化),得到的还是1x1x1000。最后对1000个通道的得分进行比较,取较大值作为预测类别。这样做的好处就是大大减少特征位置对分类带来的影响

Pytorch之VGG16图像分类_第5张图片

从上图可以看出,猫在原图片中不同的位置,如果使用全连接层很容易使得剪裁之后的图片丢失关键目标。

然而使用全卷积层,对原图片直接进行卷积,最终得分图经过求和池化后得到的得分都是1,即不管猫在图片的什么位置,都能判定这张图片中有猫,保证分类正确。保留原图,使用全卷积层,让网络模型去找猫,无论猫在那个位置,把feature map整合成一个值,如果这个值大,那么就有猫;反之,没有猫。与猫在图片中的位置无关,鲁棒性大大增强。

3.VGG创新点

1.网络层数深

VGG 实验室提出了 VGG11 、 VGG13、 VGG16 VGG19 等一系列的网络模型 ,并将网络深度最高提升至 19 层
虽然VGG层数较多,总的网络深度从11层到19层,但是它的整体结构还是相对简单。概括来说,VGG由5层卷积层(每个卷积层的子层数量不同)、3层全连接层、softmax输出层构成,层与层之间使用maxpooling(最大化池)分开,所有隐层的激活单元都采用ReLU函数。

2.小卷积核

通过堆叠多个3*3的卷积核来代替大尺度卷积核(目的:减少所需参数)。

在原论文中提到,可以通过堆叠2个3*3的卷积核替代5*5卷积核(使得2个3*3的卷积核与5*5的卷积核拥有相同的感受野)堆叠3个3*3的卷积核替代7*7的卷积核(使得3个3*3的卷积核与7*7的卷积核拥有相同的感受野)

在卷积神经网络中,决定某一层输出结果中一个元素所对应的输入层的区域大小,被称作感受野(receptive field)。通俗 的解释是,输出feature map上的一个单元 对应输入层上的区域大小。

Pytorch之VGG16图像分类_第6张图片

 如上图,最下层是一个9x9x1的特征矩阵 ,首先将其通过Conv1(大小为3x3,步距为2),通过计算公式,可以得到大小为4x4x1的特征矩阵;再将其通过最大池化下载量操作(大小为2x2,步距为2),得到一个2x2x1的大小。 

感受野计算公式:

Pytorch之VGG16图像分类_第7张图片

 Feature map(最后得到的特征图):F=1
Pool1层:其输出的是2*2大小,其输入的是4*4大小,Ksize=2,Stride=2则F=(1-1)*2+2=2
Conv1:其输出的是4*4大小,其输入的是9*9大小,Ksize=3,Stride=2则F=(2-1)*2+3=5
 

为什么会使用较小的卷积核尺寸代替具有相同感受野的大尺寸的卷积核呢?

大幅度减少模型参数数量。小卷积核选取小的 stride可以防止较大的stride导致细节信息的丢失。

②多层卷积层(每个卷积层后都有非线性激活函数),增加非线性,提升模型性能。

 例如使用7x7的卷积核所需参数,与堆叠三个3x3卷积核所需参数,假设输入输出channel为C

7x7卷积核所需参数:7x7XCXC = 49CxC

三个3x3卷积核:(3x3xCxC) x 3 =27CxC

第一个C为卷积核的通道数,第二个C是卷积核的数量。较小的卷积核参数要减少了近一半啦。

3.全卷积层

上面已经介绍了,这里就不说了。

二、VGG16实现

1.定义VGG网络模型

通过一系列的键值对,我们定义了不同VGG模型的卷积配置信息,可以以此搭建不同层数的VGG网络模型。

class VGG(nn.Module):
    def __init__(self, features, num_classes=1000, init_weights=False):
        super(VGG, self).__init__()
        self.featurs = features
        self.classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(4096, num_classes)
        )

        # 初始化参数
        if init_weights:
            self._initialize_weights()

    def forward(self, x):
        # N x 3 X 224 x224
        x = self.featurs(x)
        # N x 512 x 7 x 7
        x = torch.flatten(x, start_dim=1)
        # N x 512*7*7
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():  # 变量网络所有层
            if isinstance(m, nn.Conv2d):  # 是否为卷积层
                # 使用Kaiming初始化方法来初始化该层的权重
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:  # 否具有偏差项
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):  # 是否为Linear
                # 正太分布初始化全连接层
                nn.init.normal_(m.weight, 0, 0.01)
                # 将偏项设置为0
                nn.init.constant_(m.bias, 0)


def make_features(cfg: list):
    # 创建一个空列表用于存储神经网络的不同层
    layers = []
    # 初始输入通道数
    in_channels = 3
    # 遍历传入的配置列表
    for v in cfg:
        if v == "M":  # 池化层3
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:  # 卷积层
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            layers += [conv2d, nn.ReLU(inplace=True)]
            # 更新输入通道数,以便下一层的卷积层使用
            in_channels = v
    # 返回一个包含所有层的顺序容器,通常是一个特征提取器部分
    return nn.Sequential(*layers)


# 定义了不同VGG模型的卷积配置信息,其中 'M' 表示池化层,数字表示卷积层的输出通道数
cfgs = {
    'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'vgg19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}


# 定义一个函数 vgg,用于构建不同类型的VGG神经网络模型
def vgg(model_name="vgg16", **kwargs):
    # 检查传入的模型名是否在配置字典中
    assert model_name in cfgs, "Warning: model number {} not in cfgs dict!".format(model_name)
    # 根据模型名获取对应的卷积配置信息
    cfg = cfgs[model_name]

    # 使用 make_features 函数创建特征提取器,然后将其传递给 VGG 模型
    model = VGG(make_features(cfg), **kwargs)
    return model

2.加载数据集

这里使用花朵数据集,数据集制造和数据集使用的脚本的参考:Pytorch之AlexNet花朵分类_风间琉璃•的博客-CSDN博客

 加载数据集和测试集,并进行相应的预处理操作。

     data_transform = {
        "train": transforms.Compose([transforms.RandomResizedCrop(224),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
                                     ]),

        "val": transforms.Compose([transforms.Resize((224, 224)),
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])}

    # 数据集根目录
    data_root = os.path.abspath(os.getcwd())
    print(os.getcwd())
    # 图片目录
    image_path = os.path.join(data_root, "data_set", "flower_data")
    print(image_path)
    assert os.path.exists(image_path), "{} path does not exit.".format(image_path)

    # 准备数据集
    train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
                                         transform=data_transform["train"])
    train_num = len(train_dataset)

    validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
                                            transform=data_transform["val"])
    val_num = len(validate_dataset)


    # 定义一个包含花卉类别到索引的字典:雏菊,蒲公英,玫瑰,向日葵,郁金香
    # {'daisy':0, 'dandelion':1, 'roses':2, 'sunflower':3, 'tulips':4}
    # 获取包含训练数据集类别名称到索引的字典,这通常用于数据加载器或数据集对象中。
    flower_list = train_dataset.class_to_idx
    # 创建一个反向字典,将索引映射回类别名称
    cla_dict = dict((val, key) for key, val in flower_list.items())
    # 将字典转换为格式化的JSON字符串,每行缩进4个空格
    json_str = json.dumps(cla_dict, indent=4)
    # 打开名为 'class_indices.json' 的JSON文件,并将JSON字符串写入其中
    with open('class_indices.json', 'w') as json_file:
        json_file.write(json_str)

    batch_size = 32
    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
    print("using {} dataloader workers every process".format(nw))

    # 加载数据集
    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=batch_size, shuffle=True,
                                               num_workers=nw)

    validate_loader = torch.utils.data.DataLoader(validate_dataset,
                                                  batch_size=4, shuffle=False,
                                                  num_workers=nw)

    print("using {} images for training, {} images for validation.".format(train_num, val_num))

3.训练模型

数据集预处理完成后,就可以进行网络模型的训练了。

    model_name = "vgg16"
    net = vgg(model_name=model_name, num_classes=5, init_weights=True)
    net.to(device)

    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters(), lr=0.0001)

    epochs = 120
    best_acc = 0.0
    save_path = './{}Net.pth'.format(model_name)
    train_steps = len(train_loader)

    for epoch in range(epochs):
        # 将神经网络设置为训练模式
        net.train()
        running_loss = 0.0
        train_bar = tqdm(train_loader, file=sys.stdout)
        for step, data in enumerate(train_bar):
            images, labels = data
            optimizer.zero_grad()
            outputs = net(images.to(device))
            loss = loss_function(outputs, labels.to(device))
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            train_bar.desc = "train epoch[{}/{}] loss:{:.3f}".format(epoch + 1, epochs, loss)

4.测试模型

每训练完一个epoch,就进行一次测试,并将比较准确率大小,并根据当前模型的准确率保存模型,如果当前准确率优于之前的最佳准确率,则保存当前模型参数;否则,不变。 

        # 将神经网络设置为评估模式
        net.eval()
        acc = 0.0
        with torch.no_grad():
            val_bar = tqdm(validate_loader, file=sys.stdout)
            for val_data in val_bar:
                val_images, val_labels = val_data
                outputs = net(val_images.to(device))
                predict_y = torch.max(outputs, dim=1)[1]
                acc += torch.eq(predict_y, val_labels.to(device)).sum().item()

        val_accurate = acc / val_num
        print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' %
              (epoch + 1, running_loss / train_steps, val_accurate))

        # 如果当前模型的验证准确率优于之前的最佳准确率,则保存当前模型参数
        if val_accurate > best_acc:
            best_acc = val_accurate
            torch.save(net.state_dict(), save_path)

    print('Finished Training')

Pytorch之VGG16图像分类_第8张图片

训练完120个epoch后,基本上可以达到80%以上的准确率。 

三、实现图像分类

利用上述训练好的网络模型进行测试,验证是否能完成分类任务。

def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    data_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    # 加载图片
    img_path = 'daisy.jpg'
    assert os.path.exists(img_path), "file: '{}' does not exist.".format(img_path)
    image = Image.open(img_path)

    # img.show()
    image.show()
    # [N, C, H, W]
    img = data_transform(image)
    # 扩展维度
    img = torch.unsqueeze(img, dim=0)

    # 获取标签
    json_path = 'class_indices.json'
    assert os.path.exists(json_path), "file: '{}' does not exist.".format(json_path)
    with open(json_path, 'r') as f:
        # 使用json.load()函数加载JSON文件的内容并将其存储在一个Python字典中
        class_indict = json.load(f)

    # 加载网络
    model = vgg(model_name="vgg16", num_classes=5).to(device)

    # 加载模型文件
    weights_path = "./vgg16Net.pth"
    assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path)
    model.load_state_dict(torch.load(weights_path))

    model.eval()
    with torch.no_grad():
        # 对输入图像进行预测
        output = torch.squeeze(model(img.to(device))).cpu()
        # 对模型的输出进行 softmax 操作,将输出转换为类别概率
        predict = torch.softmax(output, dim=0)
        # 得到高概率的类别的索引
        predict_cla = torch.argmax(predict).numpy()

    res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)], predict[predict_cla].numpy())
    draw = ImageDraw.Draw(image)
    # 文本的左上角位置
    position = (10, 10)
    # fill 指定文本颜色
    draw.text(position, res, fill='red')
    image.show()
    for i in range(len(predict)):
        print("class: {:10}   prob: {:.3}".format(class_indict[str(i)], predict[i].numpy()))

运行结果:

Pytorch之VGG16图像分类_第9张图片

Pytorch之VGG16图像分类_第10张图片

结束语

感谢阅读吾之文章,今已至此次旅程之终站 。

吾望斯文献能供尔以宝贵之信息与知识也 。

学习者之途,若藏于天际之星辰,吾等皆当努力熠熠生辉,持续前行。

然而,如若斯文献有益于尔,何不以三连为礼?点赞、留言、收藏 - 此等皆以证尔对作者之支持与鼓励也 。

愿尔之学习之路风平浪静,充满希望。再次感谢尔之阅读与关注,吾期望再次相见!
 

你可能感兴趣的:(Pytorch,pytorch,人工智能,python)