图像分类篇——使用pytorch搭建VGG网络

本文为学习记录和备忘录，对代码进行了详细注释，以供学习。
内容来源：
★github： https://github.com/WZMIAOMIAO/deep-learning-for-image-processing

★b站：https://space.bilibili.com/18161609/channel/index

★CSDN：https://blog.csdn.net/qq_37541097

---

1. VGG网络详解

1.1 VGG网络概述

VGG网络在2014年由牛津大学著名研究组VGG(Visual Geometry Group)提出。
论文全称：Very Deep Convolutional Networks For Large-scale Image Recognition.

论文链接：Very Deep Convolutional Networks for Large-Scale Image Recognition

论文的亮点是：通过堆叠多个3×3的卷积核来替代大尺度卷积核，即减少所需参数。具体来讲，可以通过堆叠2个3×3的卷积核替代5×5的卷积核，堆叠3个3×3的卷积核替代7×7的卷积核，两者之间拥有相同的感受野。

VGG网络是一个Deep CNN，具备CNN所有的功能，常用来提取特征图像。在Localization和Classification tasks都取得了很大的成就。

1.2 感受野

（1）感受野的含义

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

（2）感受野的计算

计算公式：
F(i) = [F(i+1)-1]×Stride + Ksize
其中F(i)为第i层感受野，Stride为第i层的步距，Ksize为卷积核或池化核尺寸（默认Stride=1）。
下面举一个例子说明：

Feature map: F = 1；
Pool1: F = (1 - 1) x 2 + 2 = 2；
Conv1: F = (2 - 1) x 2 + 3 = 5；
即第3层1*1区域对应第2层上感受野为2x2，对应第1层上的感受野为5x5。

（3）总结

上面提到，VGG网络中，通过堆叠2个3x3的卷积核替代5x5的卷积核，堆叠3个3x3的卷积核替代7x7的卷积核。依据如下：
感受野计算公式F(i)=[F(i+1)-1]×Stride + Ksize，其中步长Stride=1，卷积核大小Ksize=3，则：
Feature map: F = 1
Conv3x3(3): F = (1 - 1) x 1 + 3 = 3
Conv3x3(2): F = (3 - 1) x 1 + 3 = 5
Conv3x3(1): F = (5 - 1) x 1 + 3 = 7
这样做的目的是减少所需参数，依据如下（假设输入输出channel为C）：使用77的卷积核所需参数为7x7xCxC = 49C²，使用3个33卷积核3x(3x3xCxC) = 27C²。（此处第一个C表示输入特征矩阵深度，第二个C表示卷积核个数即输出矩阵深度）

1.3 VGG网络结构和详细参数

作者在文中给出了6个版本：

最常用的是版本D，即VGG16（16指卷积层和全连接层共16层），其网络结构如下图所示：

• 对于Conv3：kernel_size=3；stride=1；padding=1
• 对于maxpool:size=2；stride=2

$$OUTsize=(INsize-Fsize+2P)/S+1.$$

---

2. Pytorch搭建

对于代码的解释都在注释中，方便对照查看学习。

2.1 model.py

可将VGG网络分为 提取特征网络结构(FC层之前) 和 分类网络结构(3个FC层)两个部分。
导入模块：

import torch.nn as nn
import torch

首先定义模型配置文件，在这里只搭建A、B、D、E四个版本的VGG网络模型。定义一个配置模型文件：cfgs字典。其中key代表模型，value代表具体配置，value中数字64表示卷积层中3×3卷积核的个数为64，M表示Maxpool,每一项表示一层操作。

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'],
}

有了配置文件，下一步则需要构建提取特征网络结构，定义函数make_feature，其传入的变量为配置文件的列表。

def make_features(cfg: list):  # 提取特征网络，传入变量为配置变量的list
    layers = []  # 定义一个空列表用来存放每一层的结构
    in_channels = 3  # 输入图片为rgb彩色图片，则输入通道为3
    for v in cfg:  # 通过for循环遍历配置列表，即可得到卷积层和池化层组成的一个layers列表
        if v == "M":  # 遍历配置列表，如果当前配置字符是“M”，则说明当前是最大池化层
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]  # 创建一个最大池化下采样层，下采样层的池化核大小为2，步距为2
        else:  # 如果配置字符不等于"M"，则当前为卷积层
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)  # 卷积层卷积核大小为3，padding=1,stride=1
            # 输入特征矩阵的深度为in_channel(第一层为rgb的3，其它层在下面更新),输出特征矩阵的深度为v(从配置列表中获得)
            layers += [conv2d, nn.ReLU(True)]  # 将卷积层和ReLu拼接在一起，添加到layers列表中
            in_channels = v  # 本层输出即为下一层卷积层的输入
    return nn.Sequential(*layers)
    # 通过非关键字参数(*layers)的形式传递到Sequential中。Sequential有两种传入方式：非关键字和字典

通过提取特征网络可以得到输入图片的特征矩阵，接下来需要构建分类网络结构，这里定义一个VGG类。

class VGG(nn.Module):  # 定义一个VGG类，继承nn.Module父类
    def __init__(self, features, num_classes=1000, init_weights=False):
        # 传进来由make_features生成的提取特征网络结构、所需要分类的类别个数以及是否需要对网络进行权重初始化
        super(VGG, self).__init__()
        self.features = features
        self.classifier = nn.Sequential(  # 分类网络结构
            nn.Linear(512*7*7, 4096),  # FC1，图像经过提取特征网络会生成7*7*512的特征矩阵
            nn.ReLU(True),
            nn.Dropout(p=0.5),  # Dropout函数防止过拟合
            nn.Linear(4096, 4096),  # FC2
            nn.ReLU(True),
            nn.Dropout(p=0.5),
            nn.Linear(4096, num_classes)  # FC3，输出层
        )
        if init_weights:
            self._initialize_weights()

    def forward(self, x):  # 前向传播
        # N x 3 x 224 x 224
        x = self.features(x)  # 输入图像数据x经过feature得到输出
        # N x 512 x 7 x 7
        x = torch.flatten(x, start_dim=1)  # 从第1个维度开始展平(第0个维度为batch)
        # 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):
                # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                nn.init.xavier_uniform_(m.weight)  # 如果当前网络为卷积层，采用该种方法初始化
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)  # 如果卷积核采用了偏置的话，初始化bias为0
            elif isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                # nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)  # 如果当前层为FC层，则用该方法初始化并将bias设为0

最后实例化vgg网络，上面定义的VGG类初始化需要传入3个参数

def __init__(self, features, num_classes=1000, init_weights=False):

• 实例化传入的第一个参数是features，而featuers是由make_feature函数生成的.因此传入make_feature(+所对应的配置文件)。
• **kwrags对应可变长度的字典变量，调用vgg函数所需的字典变量，可能包含所需的num_classed的值和是否需要init_weights的bool变量
  定义一个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]  # 根据所给的model_name在cfgs字典中找到相应value值作为配置list，传入下面的make_features()

    model = VGG(make_features(cfg), **kwargs)
    return model

2.2 train.py

由于VGG网络非常大，训练耗时长，如果需要使用VGG网络，最好使用迁移学习的方法来训练自己的样本集。预训练权重文件获取途径如下：

# official pretrain weights
model_urls = {
     
    'vgg11': 'https://download.pytorch.org/models/vgg11-bbd30ac9.pth',
    'vgg13': 'https://download.pytorch.org/models/vgg13-c768596a.pth',
    'vgg16': 'https://download.pytorch.org/models/vgg16-397923af.pth',
    'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth'
}

其他与AlexNet网络差别不大，主要区别在于：
①对训练集和验证集进行预处理(line20-line29)

    data_transform = {
     
        "train": transforms.Compose([transforms.RandomResizedCrop(224),  # 随机裁剪
                                     transforms.RandomHorizontalFlip(),  # 随机水平翻转
                                     transforms.ToTensor(),  # 将其转化为Tensor格式
                                     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))])}

预处理第1步有的方法将rgb3个通道减去[123.68,116.78,103.94]，即ImageNet图像数据集所有图片rgb3个通道的均值，这是迁移学习的方法，若从头开始训练则不需要，如果基于迁移学习的方法进行再训练的话，则需要进行上述操作。
②实例化VGG网络(line66-line69)

    model_name = "vgg16" 

通过model_name指定需要那一个版本的vgg配置。

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

其中后两个参数保存在model.py最下方**kwargs可变长度字典中。

训练部分全部代码如下：

import os
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


# ###########由于VGG网络非常大，训练耗时长，如果需要使用VGG网络，最好使用迁移学习的方法来训练自己的样本集###########
# #########与AlexNet网络差别不大，主要区别在于line20-line29和实例化VGG网络，line66-line69#########
def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))

    # 对训练集及验证集进行预处理操作
    data_transform = {
     
        "train": transforms.Compose([transforms.RandomResizedCrop(224),  # 随机裁剪
                                     transforms.RandomHorizontalFlip(),  # 随机水平翻转
                                     transforms.ToTensor(),  # 将其转化为Tensor格式
                                     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]),  # 标准化处理
        # 预处理第1步有的方法将rgb3个通道减去[123.68,116.78,103.94](ImageNet图像数据集所有图片rgb3个通道的均值)，这是迁移学习的方法
        # 从头开始训练则不需要，如果基于迁移学习的方法进行再训练的话，则需要进行上述操作
        "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.path.join(os.getcwd(), "../.."))  # get data root path
    image_path = os.path.join(data_root, "data_set", "flower_data")  # flower data set path
    assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
    train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"),
                                         transform=data_transform["train"])
    train_num = len(train_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())
    # write dict into json file
    json_str = json.dumps(cla_dict, indent=4)
    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_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"),
                                            transform=data_transform["val"])
    val_num = len(validate_dataset)
    validate_loader = torch.utils.data.DataLoader(validate_dataset,
                                                  batch_size=batch_size, shuffle=False,
                                                  num_workers=nw)
    print("using {} images for training, {} images for validation.".format(train_num,
                                                                           val_num))

    # test_data_iter = iter(validate_loader)
    # test_image, test_label = test_data_iter.next()

    # 实例化vgg网络
    model_name = "vgg16"   # 通过model_name指定需要那一个版本的vgg配置
    net = vgg(model_name=model_name, num_classes=5, init_weights=True)  # 后两个参数保存在model.py最下方**kwargs可变长度字典中
    net.to(device)
    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.Adam(net.parameters(), lr=0.0001)

    epochs = 30
    best_acc = 0.0
    save_path = './{}Net.pth'.format(model_name)
    train_steps = len(train_loader)
    for epoch in range(epochs):
        # train
        net.train()
        running_loss = 0.0
        train_bar = tqdm(train_loader)
        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()

            # print statistics
            running_loss += loss.item()

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

        # validate
        net.eval()
        acc = 0.0  # accumulate accurate number / epoch
        with torch.no_grad():
            val_bar = tqdm(validate_loader)
            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()

2.3 predict.py

与AlexNet网络差别不大，预测部分全部代码如下：

import os
import json

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

from model import vgg


# ########与AlexNet网络差别不大#######
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))])

    # load image
    img_path = "../tulip.jpg"
    assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
    img = Image.open(img_path)
    plt.imshow(img)
    # [N, C, H, W]
    img = data_transform(img)
    # expand batch dimension
    img = torch.unsqueeze(img, dim=0)

    # read class_indict
    json_path = './class_indices.json'
    assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)

    json_file = open(json_path, "r")
    class_indict = json.load(json_file)
    
    # create model
    model = vgg(model_name="vgg16", num_classes=5).to(device)
    # load model weights
    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, map_location=device))

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

    print_res = "class: {}   prob: {:.3}".format(class_indict[str(predict_cla)],
                                                 predict[predict_cla].numpy())
    plt.title(print_res)
    print(print_res)
    plt.show()


if __name__ == '__main__':
    main()