基于VGG16 实现分类问题的训练与预测（pytorch版本）

MyVGG.py

import torch
from torchvision import models
import torch.nn as nn

vgg16 = models.vgg16(pretrained=True)
vgg = vgg16.features
for param in vgg.parameters():
    param.requires_grad_(False)


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, 5),
            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

EarlyStopping.py

import numpy as np
import torch

class EarlyStopping:
    """Early stops the training if validation loss doesn't improve after a given patience."""

    def __init__(self, patience=7, verbose=True, delta=0):
        """
        Args:
            patience (int): How long to wait after last time validation loss improved.
                            Default: 7
            verbose (bool): If True, prints a message for each validation loss improvement.
                            Default: False
            delta (float): Minimum change in the monitored quantity to qualify as an improvement.
                            Default: 0
        """
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.val_loss_min = np.Inf
        self.delta = delta

    def __call__(self, val_loss, model):

        score = -val_loss

        if self.best_score is None:
            self.best_score = score
            self.save_checkpoint(val_loss, model)
        elif score < self.best_score + self.delta:
            self.counter += 1
            print(
                'EarlyStopping counter: {self.counter} out of {self.patience}')
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.save_checkpoint(val_loss, model)
            self.counter = 0
        return self.early_stop

    def save_checkpoint(self, val_loss, model):
        '''Saves model when validation loss decrease.'''
        if self.verbose:
            print(
                f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')

        torch.save(model, "data/trained/VGG_trained.pkl")
        self.val_loss_min = val_loss

train.py

import torch
import torch.nn as nn
import torch.utils.data as Data
from torchvision import transforms
from torchvision.datasets import ImageFolder
import hiddenlayer as hl

from EarlyStopping import EarlyStopping
from MyVGG import MyVggModel

# 定义transform
data_transforms = transforms.Compose([
    transforms.RandomResizedCrop(224),# 随机长宽比裁剪为224*224
    transforms.RandomHorizontalFlip(),# 依概率p=0.5水平翻转
    transforms.ToTensor(), # 转化为张量并归一化至[0-1]
    ## 图像标准化处理
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# 加载数据，划分训练集与测试集
full_data_dir = "data/chap6/training"
full_data = ImageFolder(full_data_dir,transform=data_transforms)
train_size = int(0.8 * len(full_data_dir))
val_size = len(full_data) - train_size
train_dataset, val_dataset = torch.utils.data.random_split(full_data, [train_size, val_size])
train_dataset=train_dataset.dataset
val_dataset=val_dataset.dataset

train_data_loader = Data.DataLoader(train_dataset,batch_size=32,
                                  shuffle=True,num_workers=2)
val_data_loader = Data.DataLoader(val_dataset,batch_size=32,
                                  shuffle=True,num_workers=2)
# 输出样本数
print("训练集样本数:",len(train_dataset.targets))
print("验证集样本数:",len(val_dataset.targets))

# 网络和early_stop
Myvggc = MyVggModel().cuda()
early_stopping = EarlyStopping(patience=15, verbose=True)


# 定义优化器，损失函数
optimizer = torch.optim.Adam(Myvggc.parameters(), lr=0.0003)
loss_func = nn.CrossEntropyLoss()

# 记录训练过程的指标，使用Canvas进行可视化
history1 = hl.History()
canvas1 = hl.Canvas()

# 对模型进行迭代训练,对所有的数据训练EPOCH轮
for epoch in range(2000):
    train_loss_epoch = 0
    val_loss_epoch = 0
    train_corrects =0
    val_corrects = 0

    # 对训练数据的迭代器进行迭代计算
    Myvggc.train()
    for step, (b_x, b_y) in enumerate(train_data_loader):
        b_x, b_y = b_x.cuda(), b_y.cuda()
        output = Myvggc(b_x)            # CNN在训练batch上的输出
        loss = loss_func(output, b_y)   # 交叉熵损失函数
        pre_lab = torch.argmax(output,1)
        optimizer.zero_grad()           # 每个迭代步的梯度初始化为0
        loss.backward()                 # 损失的后向传播，计算梯度
        optimizer.step()                # 使用梯度进行优化
        train_loss_epoch += loss.item() * b_x.size(0)
        train_corrects += torch.sum(pre_lab == b_y.data)

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

    # 计算在验证集上的表现
    Myvggc.eval()
    for step, (val_x, val_y) in enumerate(val_data_loader):
        val_x, val_y = val_x.cuda(), val_y.cuda()
        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_dataset.targets)
    val_acc = val_corrects.double() / len(val_dataset.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"])
        canvas1.draw_plot(history1["train_acc"])

    if early_stopping(train_loss, Myvggc) is True:
        break

# 保存模型
torch.save(Myvggc, "data/trained/VGG_trained.pkl")

predict.py

import numpy as np
import torch
import cv2
import torch.utils.data as Data

from torchvision import transforms
from torchvision.datasets import ImageFolder

Myvggc2 = torch.load("/home/xcy/torch/data/trained/VGG_trained.pkl")
Myvggc2 = Myvggc2.cuda()

test_data_transforms = transforms.Compose([
    transforms.ToTensor(),  # 转化为张量并归一化至[0-1]
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

test_data_dir = "/home/xcy/test"
test_data = ImageFolder(test_data_dir, transform=test_data_transforms)
test_data_loader = Data.DataLoader(test_data, batch_size=32,
                                  shuffle=False, num_workers=2)

print("测试集样本数:", len(test_data.targets))

test_corrects = 0

Myvggc2.eval()

with torch.no_grad():
    for step, (test_x, test_y) in enumerate(test_data_loader):
        test_x, test_y = test_x.cuda(), test_y.cuda()
        print(test_x.shape)
        output = Myvggc2(test_x)
        print(output)
        pre_lab = torch.argmax(output, 1)
        print("predict as: ", pre_lab)
        test_corrects += torch.sum(pre_lab == test_y.data+1)
        test_acc = test_corrects.double() / len(test_data.targets)
        print(test_acc)