微调预训练的VGG16网络

下面展示一些 内联代码片。

微调预训练的VGG16网络

import numpy as np
import hiddenlayer as hl
import torch
import torch.nn as nn
import torch.utils.data as Data
from torchvision import models
from torchvision import transforms
from torchvision.datasets import ImageFolder
import matplotlib.pyplot as plt


vgg16 = models.vgg16(pretrained=True)  # 导入预训练好的vgg16模型
vgg = vgg16.features  # 获取vgg16的特征提取层
# 将vgg16的特征提取层参数冻结，不对其进行更新
for param in vgg.parameters():
    param.requires_grad_(False)


# vgg16的特征提取层 + 新的全连接层    组成新的网络
class MyVggModel(nn.Module):
    def __init__(self):
        super(MyVggModel, self).__init__()
        # 预训练的vgg16的特征提取层
        self.vgg = vgg
        # 添加新的全连接层
        self.classifier = nn.Sequential(
            # 第一个全连接层
            nn.Linear(25088, 512),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            # 第二个全连接层
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(p=0.5),
            # 第三个全连接层
            nn.Linear(256, 10),
            nn.Softmax(dim=1)
        )

    def forward(self, x):
        x = self.vgg(x)
        x = x.view(x.size(0), -1)
        output = self.classifier(x)
        return output


# 准备数据集
# 使用10类猴子的数据集，对训练集预处理
train_data_transforms = transforms.Compose(
    [
        transforms.RandomResizedCrop(244),  # 随机长宽比裁剪为224×224
        transforms.RandomHorizontalFlip(),  # 依概率p=0.5水平翻转
        transforms.ToTensor(),  # 转化为张量
        # 将图像标准化处理
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]
)
# 对验证集的预处理
val_data_transforms = transforms.Compose(
    [
        transforms.Resize(256),  # 重置图像分辨率
        transforms.CenterCrop(224),  # 依据给定的size从中心裁剪
        transforms.ToTensor(),  # 转化为张量
        # 将图像标准化处理
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]
)


# 读取图像
# 读取训练图像
train_data_dir = "./data/10-monkey-species/training"
train_data = ImageFolder(train_data_dir, transform=train_data_transforms)
# shuffle=True:打乱数据集    num_workers=2：一次性创建2个工作进程
train_data_loader = Data.DataLoader(train_data, batch_size=32, shuffle=True)
# 读取验证集
val_data_dir = "./data/10-monkey-species/validation"
val_data = ImageFolder(val_data_dir, transform=val_data_transforms)
# shuffle=True:打乱数据集    num_workers=2：一次性创建2个工作进程
val_data_loader = Data.DataLoader(val_data, batch_size=32, shuffle=True)


# 显示图像
# 获得一个batch的图像
for step, (b_x, b_y) in enumerate(train_data_loader):
    if step > 0:
        break

    # 显示一个batch的图像
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    plt.figure(figsize=(12, 6))
    for ii in np.arange(len(b_y)):
        plt.subplot(4, 8, ii+1)   # 分割成4行8列，显示在ii+1个位置
        image = b_x[ii, :, :, :].numpy().transpose((1, 2, 0))
        image = image*std + mean
        image = np.clip(image, 0, 1)
        plt.imshow(image)
        plt.title(b_y[ii].data.numpy())
        plt.axis("off")
    plt.subplots_adjust(hspace=0.3)  # 调整子图布局
    plt.show()


if __name__ == "__main__":
    Myvggc = MyVggModel()
    print(Myvggc)
    # 定义优化器
    optimize = torch.optim.Adam(Myvggc.parameters(), lr=0.003)
    # 定义损失函数
    loss_func = nn.CrossEntropyLoss()
    # 记录训练过程的指标
    history1 = hl.History()
    # 使用Canvas进行可视化
    canvas1 = hl.Canvas()

    # 对模型进行训练，所有的数据训练epoch轮
    for epoch in range(10):
        train_loss_epoch = 0
        val_loss_epoch = 0
        train_corrects = 0
        val_corrects = 0

        Myvggc.training()
        for step, (b_x, b_y) in enumerate(train_data_loader):
            # 计算每个batch的损失和精度
            # 向前传播
            output = Myvggc(b_x)
            loss = loss_func(output, b_y)
            pre_lab = torch.argmax(output, 1)
            # 向后传播
            optimize.zero_grad()
            loss.backward()
            optimize.step()
            train_loss_epoch += loss.item() * b_x.size(0)
            train_corrects += torch.sum(pre_lab == b_y.data)

            print('[Epoch:%d, step:%d]' % (epoch, step))

        # 计算一个epoch的损失和精度
        train_loss = train_loss_epoch / len(train_data.targets)
        train_acc = train_corrects.double() / len(train_data.targets)

        # 计算在验证集上的表现
        Myvggc.eval()
        for step, (val_x, val_y) in enumerate(val_data_loader):
            output = Myvggc(val_x)
            loss = loss_func(output, val_y)
            pre_lab = torch.argmax(output, 1)
            val_loss_epoch += loss.item() * val_x.size(0)
            val_corrects += torch.sum(pre_lab == val_y.data)

        # 计算一个epoch的损失和精度
        val_loss = val_loss_epoch / len(val_data.targets)
        val_acc = val_corrects.double() / len(val_data.targets)

        # 保存每个epoch上的输出loss和acc
        history1.log(epoch,
                     train_loss=train_loss,
                     val_loss=val_loss,
                     train_acc=train_acc.item(),
                     val_acc=val_acc.item()
                     )

        # 可视化网络训练过程
        with canvas1:
            canvas1.draw_plot([history1["train_loss"], history1["val_loss"]])
            canvas1.draw_plot([history1["train_acc"], history1["val_acc"]])