文章目录
- Pytorch模型自定义数据集训练流程
-
- 1、任务描述
- 2、导入各种需要用到的包
- 3、分割数据集
- 4、将数据转成pytorch标准的DataLoader输入格式
- 5、导入预训练模型，并修改分类层
- 6、开始模型训练
- 7、利用训好的模型做预测
Pytorch模型自定义数据集训练流程
我们以kaggle竞赛中的猫狗大战数据集为例搭建Pytorch自定义数据集模型训练的完整流程。
1、任务描述
Cats vs. Dogs（猫狗大战）数据集是Kaggle大数据竞赛某一年的一道赛题，利用给定的数据集，用算法实现猫和狗的识别。 其中包含了训练集和测试集，训练集中猫和狗的图片数量都是12500张且按顺序排序，测试集中猫和狗混合乱序图片一共12500张。
下载地址：https://www.kaggle.com/c/dogs-vs-cats/data
卷积神经网络（CNN）是一类包含卷积计算且具有深度结构的前馈神经网络，是深度学习的代表算法之一。卷积神经网络具有表征学习能力，能够按其阶层结构对输入信息进行平移不变分类，因此也被称为“平移不变人工神经网络”。
默认对图像分类各种算法已经熟悉，卷积、池化、批量归一化、全连接等各种结构具体细节这里不讨论，有不懂的可自行学习。
2、导入各种需要用到的包
import torch
import torchvision
from torchvision import datasets, transforms
import torch.utils.data
import matplotlib.pyplot as plt
from torch.utils.data import TensorDataset,DataLoader,Dataset
from torch.utils.tensorboard import SummaryWriter
import torch.nn.functional as F
from torch import nn
import numpy as np
import os
import shutil
from PIL import Image
import warnings
warnings.filterwarnings("ignore")
3、分割数据集
下载猫狗大战数据集，并解压。
解压完成后，通过以下代码实现数据集预处理（剔除不能正常打开的图片，打乱数据集等）；然后对数据集进行分割，其中90%的数据集作为train训练，另外10%的数据集作为test测试。
# 分割数据集，将全部数据分成0.9的Train和0.1的Test
source_path = r"./kagglecatsanddogs_5340/PetImages/"
# 如果不存在文件夹要新建一个
if not os.path.exists(os.path.join(source_path, "train")):
os.mkdir(os.path.join(source_path, "train"))
train_dir = os.path.join(source_path, "train")
if not os.path.exists(os.path.join(source_path, "test")):
os.mkdir(os.path.join(source_path, "test"))
test_dir = os.path.join(source_path,"test")
## 将Cat和Dog文件夹全部移到train目录下，然后再从train目录下移动10%到test目录下
for category_dir in os.listdir(source_path):
if category_dir not in ["train", "test"]:
shutil.move(os.path.join(source_path,category_dir), os.path.join(source_path,"train"))
## 开始移动，移动前先剔除不能正常打开的图片
for dir in os.listdir(train_dir):
category_dir_path = os.path.join(train_dir, dir)
image_file_list = os.listdir(category_dir_path) # 取出全部图片文件
for file in image_file_list:
try:
Image.open(os.path.join(category_dir_path, file))
except:
os.remove(os.path.join(category_dir_path, file))
image_file_list.remove(file)
np.random.shuffle(image_file_list)
test_num = int(0.1*len(image_file_list))
#移动10%文件到对应目录
if not os.path.exists(os.path.join(test_dir,dir)):
os.mkdir(os.path.join(test_dir,dir))
if len(os.listdir(os.path.join(test_dir,dir))) < test_num: # 只有未移动过才需要移动，否则每运行一次都会移动一下
for i in range(test_num):
shutil.move(os.path.join(category_dir_path,image_file_list[i]), os.path.join(test_dir,dir,image_file_list[i]))
4、将数据转成pytorch标准的DataLoader输入格式
1、先对数据集进行预处理，包括resize成224*224的尺寸，因为vgg_net模型需要的输入尺寸为[N, 224, 224, 3]；随机翻转，随机旋转等，另外对数据集做Normalize标准化，其中的mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.2]是从ImageNet数据集上的百万张图片中随机抽样计算得到的，以上这些内容主要是数据增强，增强模型的泛化性，有更好的预测效果。
2、然后将预处理好的数据转成pytorch标准的DataLoader输入格式，。
# 数据预处理
transform = transforms.Compose([
transforms.RandomResizedCrop(224),# 对图像进行随机的crop以后再resize成固定大小
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.2]), # ImageNet全部图片的平均值和标准差
transforms.RandomRotation(20), # 随机旋转角度
transforms.RandomHorizontalFlip(p=0.5), # 随机水平翻转
])
# 读取数据
root = source_path
train_dataset = datasets.ImageFolder(root + '/train', transform)
test_dataset = datasets.ImageFolder(root + '/test', transform)
# 导入数据
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=16, shuffle=False)
5、导入预训练模型，并修改分类层
1、定义device，如果有GPU模型训练会自动用GPU训练，否则会使用CPU；使用GPU训练，只需在模型、数据、损失函数上使用cuda()就行。
2、这边默认对分类图像算法都熟悉，可以自己构建vgg16的完整网络，在猫狗数据集上重新训练。也可以下载预训练模型，由于原网络的分类输出是1000类别的，但是我们的图片只有两类，所以需要修改分类层，让模型能够适配我们的训练数据集。
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
vgg16 = torchvision.models.vgg16(pretrained=True).to(device)
print(vgg16)
inputs = torch.rand(1, 3, 224, 224) # 拿一个随机tensor测试一下网络的输出是否满足预期
output = vgg16(inputs.to(device))
print("原始VGG网络的输出：",output.size())
# 构建新的全连接层
vgg16.classifier = torch.nn.Sequential(torch.nn.Linear(25088, 100),
torch.nn.ReLU(),
torch.nn.Dropout(p=0.5),
torch.nn.Linear(100, 2)).to(device)
inputs = torch.rand(1, 3, 224, 224)
output = vgg16(inputs.to(device))
print("新构建的VGG网络的输出：",output.size())
6、开始模型训练
开始模型训练，我们这里只训练全连接分类层，将特征层的梯度requires_grad设置为False，特征层的参数将不参与训练。
训练过程中保存效果最好的网络模型，以防掉线，可以从断点开始继续训练，同时也可以用来做预测。
训练完成后，保存训练好的网络和参数，后面可以加载模型做预测。
writer = SummaryWriter("./logs/model")
loss_func = nn.CrossEntropyLoss().to(device)
learning_rate = 0.0001
#如果我们想只训练模型的全连接层
for param in vgg16.features.parameters():
param.requires_grad = False
optimizer = torch.optim.Adam(vgg16.parameters(),lr=learning_rate)
## 断点续训开始
resume = False
if resume:
# 恢复上次的训练状态
print("Resume from checkpoint...")
checkpoint = torch.load("./models/checkpoint/ckpt_best.pth")
vgg16.load_state_dict(checkpoint['net'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch_ = checkpoint['epoch'] + 1
#从上次记录的损失和正确率接着记录
loss = checkpoint['loss']
total_test_loss = checkpoint["total_test_loss"]
total_acc = checkpoint["total_acc"]
else:
total_acc = 0.0
epoch_ = 0
##训练开始
total_train_step = 0
total_test_step = 0
min_acc = 0.0
for epoch in range(epoch_ , 10):
print("-----------train epoch {} start---------------".format(epoch))
vgg16.train()
for data in train_loader:
optimizer.zero_grad()
img, label = data
output = vgg16(img.to(device))
loss = loss_func(output, label.to(device))
loss.backward()
optimizer.step()
total_train_step += 1
if total_train_step % 10 == 0:
print("steps: {}, train_loss: {}".format(total_train_step, loss.item()))
writer.add_scalar("train_loss", loss.item(), total_train_step)
## 测试开始，看训练效果是否满足预期
total_test_loss = 0
vgg16.eval()
with torch.no_grad():
for data in test_loader:
optimizer.zero_grad()
img, label = data
output = vgg16(img.to(device))
loss = loss_func(output, label.to(device))
total_test_loss += loss
accuary = torch.sum(output.argmax(1) == label.to(device))
total_acc += accuary
total_test_step += 1
val_acc = total_acc.item() / len(test_dataset)
total_acc = 0.0
## 保存Acc最小的模型
if val_acc > min_acc:
min_acc = val_acc
torch.save(vgg16.state_dict(), "./models/2classes_vgg16_weight.pth")
print("测试Acc: {} \n 模型保存成功！".format(min_acc))
# 保存模型和训练参数的全相关信息，方便断点续训
checkpoint = {
"net": vgg16.state_dict(),
'optimizer':optimizer.state_dict(),
"loss": loss,
"epoch": epoch,
"total_test_loss": total_test_loss,
"total_acc": total_acc
}
if not os.path.exists("./models/checkpoint"):
os.mkdir("./models/checkpoint")
torch.save(checkpoint, './models/checkpoint/ckpt_best.pth')
print("测试loss: {}".format(total_test_loss.item()))
print("测试Acc: {}".format(val_acc))
writer.add_scalar("test_loss", total_test_loss.item(), total_test_step)
writer.add_scalar("test_Acc", val_acc, total_test_step)
torch.save(vgg16.state_dict(), "./models/2classes_vgg16_latest_{}.pth".format(val_acc))
7、利用训好的模型做预测
拿出一张图片做预测，首先导入预训练模型，同样改掉分类层，然后导入预训练权重，预测图片类别，输出标签值和预测类别。
import matplotlib.pyplot as plt
img_path = r"./kagglecatsanddogs_5340/PetImages/test/Cat/1381.jpg" # 拿出要预测的图片
image = Image.open(img_path).convert("RGB")
image.show()
vgg16_pred = torchvision.models.vgg16(pretrained=True)
vgg16_pred.classifier = torch.nn.Sequential(torch.nn.Linear(25088, 100),
torch.nn.ReLU(),
torch.nn.Dropout(p=0.5),
torch.nn.Linear(100, 2))
transform = torchvision.transforms.Compose([
torchvision.transforms.Resize((224,224), interpolation=2),
torchvision.transforms.ToTensor()
])
vgg16_pred.load_state_dict(torch.load("./models/2classes_vgg16_weight_15_0.9467513434294089.pth", map_location=torch.device('cpu')))
print(vgg16_pred)
image = transform(image)
print(image.size())
image = torch.reshape(image, [1,3,224,224])
vgg16_pred.eval()
with torch.no_grad():
output = vgg16_pred(image)
# print("预测值为：",output)
print("预测标签为：",output.argmax(1).item())
print("预测动物为：",train_dataset.classes[output.argmax(1)])