Pytorch自定义数据集模型完整训练流程

文章目录

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/"# 如果不存在文件夹要新建一个ifnot 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")ifnot 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 notin["train","test"]:
        shutil.move(os.path.join(source_path,category_dir), os.path.join(source_path,"train"))## 开始移动，移动前先剔除不能正常打开的图片fordirin os.listdir(train_dir):
    category_dir_path = os.path.join(train_dir,dir)
    image_file_list = os.listdir(category_dir_path)# 取出全部图片文件forfilein 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%文件到对应目录ifnot os.path.exists(os.path.join(test_dir,dir)):
        os.mkdir(os.path.join(test_dir,dir))iflen(os.listdir(os.path.join(test_dir,dir)))< test_num:# 只有未移动过才需要移动，否则每运行一次都会移动一下for i inrange(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 =Falseif 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.0for epoch inrange(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 +=1if 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
    }ifnot 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)])