VGG网络实现

一：model.py

import torch.nn as nn
import torch

# 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'
}


class VGG(nn.Module):
    def __init__(self, features, num_classes=1000, init_weights=False):
        super(VGG, self).__init__()
        self.features = features         # make_features所定义的特征网络层结构
        self.classifier = nn.Sequential(    # 分类网络层结构
            nn.Linear(512*7*7, 2048),
            nn.ReLU(True),
            nn.Dropout(p=0.5),
            nn.Linear(2048, 2048),
            nn.ReLU(True),
            nn.Dropout(p=0.5),
            nn.Linear(2048, num_classes)
        )
        # 是否进行初始化权重
        if init_weights:
            self._initialize_weights()

    ######## 网络的前向传播
    def forward(self, x):
        # N x 3 x 224 x 224
        x = self.features(x)     # 特征网络结构
        # N x 512 x 7 x 7
        x = torch.flatten(x, start_dim=1)    # 展平处理 第0个维度为batch 所以从第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):    # 卷积层所采用的操作
                # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:     # 卷积核采用了偏置则重置为0
                    nn.init.constant_(m.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)

###### 提取特征网络结构
def make_features(cfg: list):   # 传入对应网络配置的列表
    layers = []
    in_channels = 3
    for v in 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(True)]
            in_channels = v   # 下一层的输入变为这一层的卷积核个数
    return nn.Sequential(*layers)  # 将所有的层结构按顺序包裹

    # 网络配置
cfgs = {    # 第一层64个卷积核 第二层为最大池化下采样层
    '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网络
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]

            # def __init__(self, features, num_classes=1000, init_weights=False):
                  # 传入的参数第一个为features
    model = VGG(make_features(cfg), **kwargs)
    return model

# vgg_model = vgg(model_name='vgg16')

二：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("using {} device.".format(device))

# ----------------imageNET的RGB三通道的均值： [123.68,116.78,103.94]  基于迁移学习的话要减去这三个值

    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))])}
                                            # 获取当前文件路径
    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)
    # 将分类的字典编码成一个json文件

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

    # -------------- 实例化网络
    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 = 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()
        # 控制是否使用dropout  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))   # 数据加载到GPU去运行
            loss = loss_function(outputs, labels.to(device))   # 计算损失函数
            loss.backward()   # 反向传播
            optimizer.step()   # 更新节点参数

            # print statistics
            running_loss += loss.item()    # loss值累加

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

        # validate
        net.eval()
        # 控制是否使用dropout  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，则保存为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))])

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

    with open(json_path, "r") as f:
        class_indict = json.load(f)
    
    # 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)
    for i in range(len(predict)):
        print("class: {:10}   prob: {:.3}".format(class_indict[str(i)],
                                                  predict[i].numpy()))
    plt.show()


if __name__ == '__main__':
    main()