import hiddenlayer as hl
import torch
import torch.nn as nn
from torch.optim import SGD, Adam
import torch.utils.data as Data
from torchvision import models
from torchvision import transforms
from torchvision.datasets import ImageFolder
# 导入预训练好的vgg16网络
vgg16 = models.vgg16(pretrained=True)
# 获取vgg16特征提取层
vgg = vgg16.features
# 将vgg16的特征提取层参数冻结不对其进行更新
for param in vgg.parameters():
param.requires_grad_(False)
# 使用vgg16的特征提取层+其他的层自己创建新的网络
class MyVggModel(nn.Module):
def __init__(self):
super(MyVggModel, self).__init__() # 对继承自父类的属性进行初始化
self.vgg = vgg # 调用vgg16网络
# 添加新的全连接层
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
if __name__ == "__main__":
# 实例化一个网络
Myvggc = MyVggModel()
# print(Myvggc)
# 使用10类猴子的数据集,对训练集预处理----将各种操作用compose集合起来
train_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])
])
# 对验证集进行预处理
val_data_transfronms = transforms.Compose([
transforms.Resize(256), # 重置图像分辨率
transforms.CenterCrop(224), # 依据给定的大小从中间裁剪
transforms.ToTensor(),# 转化为张量并归一化
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) # 图像标准化处理
])
#加载训练数据,使用对应的数据加载器
train_data_dir = "data/monkey-species/training/training"
train_data = ImageFolder(train_data_dir, transform=train_data_transforms)
train_data_loader = Data.DataLoader(train_data, batch_size=32, shuffle=True, num_workers=2)
# 加载验证数据集,使用对应的数据加载器
val_data_dir = "data/monkey-species/validation/validation"
val_data = ImageFolder(val_data_dir, transform=val_data_transfronms)
val_data_loader = Data.DataLoader(val_data, batch_size=32, shuffle=True, num_workers=2)
print("训练集样本:", train_data)
print("训练集样本数:", len(train_data.targets))
print("验证集样本数:", len(val_data.targets))
# for step, (b_x, b_y) in enumerate(train_data_loader):
# if step > 0:
# break
#
# mean = np.array([0.485, 0.456, 0.406])
# std = np.array([0.229, 0.224, 0.225])
# plt.figure(figsize=(12, 6))
# print("b_y", b_y)
# for ii in np.arange(len(b_y)):
# print("len(b_y)", len(b_y))
# print("ii:", ii)
# plt.subplot(4, 8, ii + 1)
# image = b_x[ii, :, :, :].numpy().transpose((1, 2, 0))
# image = std * image + 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()
# 定义优化器和损失函数
optimizer = 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):
# 每个epoch前初始化数据
train_loss_epoch = 0
val_loss_epoch = 0
train_correct = 0
val_correct = 0
# 对训练数据的加载器进行迭代训练
Myvggc.train()
# 对每个batch的数据进行训练
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)
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss_epoch += loss.item() * b_x.size(0)
train_correct += torch.sum(pre_lab == b_y.data)
# 计算一个epoch的损失和精度
train_loss = train_loss_epoch / len(train_data.targets)
train_acc = train_correct.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_correct += torch.sum(pre_lab == val_y.data)
# 计算一个batch 的损失和精度
val_loss = val_loss_epoch / len(val_data.targets)
val_acc = val_correct.double() / len(val_data.targets)
# 保存每个batch上的损失和精度
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"]])