【深度学习-图像分类】PyTorch小白大战VGGNet

写在前面:

        本文的侧重点在于Pytorch实战,对于网络的理论部分不做过多的介绍。

一、VGGNet结构:

        “一句话概括VGGNet的网络结构:一卷到底”

        本文中所使用的是VGGNet-16,包含了16个卷积层。

 

        VGGNet是从AlexNet进化而来。相比于AlexNet而言,具有更小的卷积核,都是3x3的,而Alex-net卷积核较大(11x11,7x7,5x5)。并且相比于AlexNet的3x3的池化核,VGG全部为2x2的池化核。

 二、网络搭建

         知晓VGGNet结构之后,可以着手网络的搭建了。

        新建module.py脚本,编写VGGNet类,在其中编写3个函数,分别描述网络结构,正向传播过程和网络初始化过程。此外,采用列表的方法构建vgg16网络,十分的方便。

import torch.nn as nn
import torch


# 定义VGG的类,进行模型的搭建。
class VGG(nn.Module):
    def __init__(self, features, num_classes=1000, init_weights=False):
        super(VGG, self).__init__()
        self.features = 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):               # 定义前向传播过程
        # 输入图像尺寸:3×224×224
        x = self.features(x)
        # 输出图像尺寸:512×7×7
        x = torch.flatten(x, start_dim=1)
        # 展平操作
        x = self.classifier(x)
        return x

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                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):
                    nn.init.normal_(m.weight, 0, 0.01)  # 均值为0,标准差为0.01
                    nn.init.constant_(m.bias, 0)  # 偏置为0


# 利用列表去构建网络
def make_features(cfg: list):
    layers = []         # 构建一个空列表
    in_channels = 3
    for v in cfg:       # 遍历cfg文件中的元素
        if v == "M":
            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)


# vgg16网络结构列表。
vgg16 = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']


def vgg(**kwargs):
    model = VGG(make_features(vgg16), **kwargs)
    return model

三、网络训练

         网络搭建好之后,就可以着手对网络进行训练了。训练所使用的数据集为花分类数据集,下载地址如下所示。

https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz

        训练之前我们需要做好以下的准备工作:

        1、创建data_set文件夹,在此文件夹下创建新文件夹"flower_data"。

        2、下载好花分类数据集,并解压到创建的flower_data文件夹下面。

        3、将数据集划分成训练集train和验证集val。

        文件结构如下所示:

 训练的方法是通过编写train.py脚本来实现,训练结果打印在终端中。

import os
import sys
import json

import torch
import torch.nn as nn
from torchvision import transforms, datasets
import torch.optim as optim
from tqdm import tqdm

from model import vgg


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"using {device} device.")

    #   数据预处理的字典文件
    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))])
    }

    # 设定数据路径
    image_path = './data_set/flower_data'
    # 路径检查
    assert os.path.exists(image_path), f"{image_path} path does not exist."
    # 训练集
    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文件中
    json_str = json.dumps(cla_dict, indent=4)
    with open('class_indices.json', 'w') as json_file:
        json_file.write(json_str)

    batch_size = 8
    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])
    print(f"Using {nw} dataloader workers every process")
    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=batch_size,
                                                  shuffle=False,
                                                  num_workers=nw)
    # 打印训练集和验证集的信息。
    print(f"Using {train_num} images for training, {val_num} images.")

    # 将模型映射到设备中
    net = vgg(num_classes=5, init_weights=True)
    net.to(device)
    loss_function = nn.CrossEntropyLoss()       # 使用交叉熵损失
    optimizer = optim.Adam(net.parameters(), lr=0.0001)

    epochs = 30
    best_acc = 0.0          # float
    save_path = './vgg16Net.pth'
    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)

        # 验证集
        net.eval()
        acc = 0.0       # accumulate accurate number / epoch
        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")


if __name__ =='__main__':
    main()

四、模型验证

        网络训练好之后,需要对网络的训练结果进行检验。方法是编写predict.py脚本进行验证。在终端中打印出图像分类后的信息。

import os
import json

import torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt

from model import vgg


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))])

    # 图像载入
    for i in range(5):
        img_path = './test_pictures/{}.jpg'.format(i)
        assert os.path.exists(img_path), "file: '{}' does not exist.".format(img_path)  # 对路径是否存在进行检查
        img = Image.open(img_path)      # 将图片进行打开
        img = data_transform(img)           # 将图片进行预处理
        img = torch.unsqueeze(img, dim=0)   # 进行维度扩大,在第一个参数上进行扩长

        # 载入json文件
        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:
            class_indict = json.load(f)

        # 将路径载入到模型中,首先要创建模型
        model = vgg(num_classes=5).to(device)
        # 载入模型权重
        weights_path = "./vgg16Net.pth"
        assert os.path.exists(weights_path), f"file: '{weights_path}' does not exist."
        model.load_state_dict(torch.load(weights_path, map_location=device))

        model.eval()
        with torch.no_grad():
            # 类别预测
            output = torch.squeeze(model(img.to(device))).cpu()
            predict = torch.softmax(output, dim=0)
            predict_cla = torch.argmax(predict).numpy()

        print_res = "class: {}  prob: {:.3f}".format(class_indict[str(predict_cla)],
                                                     predict[predict_cla].numpy())
        plt.title(print_res)
        for x in range(len(predict)):
            print("class: {:10}  prob: {:.3}".format(class_indict[str(x)],
                                                     predict[x].numpy()))


if __name__ == '__main__':
    main()

        实际测试中可以发现,效果比AlexNet好不少。

 写在后面的话:

          作为一个刚刚涉及深度学习的小白,对于Python和Pytorch的掌握还是很不熟练,所以有些代码可能看起来非常的笨拙甚至于可笑,希望大家能够多多包涵。

        在此感谢UP主:@霹雳吧啦Wz。本文的代码主要在其开源的代码的基础上进行修改,再次表示感谢

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值