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)
