Pytorch模型自定義數(shù)據(jù)集訓(xùn)練流程

我們以kaggle競賽中的貓狗大戰(zhàn)數(shù)據(jù)集為例搭建Pytorch自定義數(shù)據(jù)集模型訓(xùn)練的完整流程。

1、任務(wù)描述

Cats vs. Dogs（貓狗大戰(zhàn)）數(shù)據(jù)集是Kaggle大數(shù)據(jù)競賽某一年的一道賽題，利用給定的數(shù)據(jù)集，用算法實(shí)現(xiàn)貓和狗的識別。 其中包含了訓(xùn)練集和測試集，訓(xùn)練集中貓和狗的圖片數(shù)量都是12500張且按順序排序，測試集中貓和狗混合亂序圖片一共12500張。
下載地址：https://www.kaggle.com/c/dogs-vs-cats/data

卷積神經(jīng)網(wǎng)絡(luò)（CNN）是一類包含卷積計(jì)算且具有深度結(jié)構(gòu)的前饋神經(jīng)網(wǎng)絡(luò)，是深度學(xué)習(xí)的代表算法之一。卷積神經(jīng)網(wǎng)絡(luò)具有表征學(xué)習(xí)能力，能夠按其階層結(jié)構(gòu)對輸入信息進(jìn)行平移不變分類，因此也被稱為“平移不變?nèi)斯ど窠?jīng)網(wǎng)絡(luò)”。
默認(rèn)對圖像分類各種算法已經(jīng)熟悉，卷積、池化、批量歸一化、全連接等各種結(jié)構(gòu)具體細(xì)節(jié)這里不討論，有不懂的可自行學(xué)習(xí)。

2、導(dǎo)入各種需要用到的包

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、分割數(shù)據(jù)集

下載貓狗大戰(zhàn)數(shù)據(jù)集，并解壓。
解壓完成后，通過以下代碼實(shí)現(xiàn)數(shù)據(jù)集預(yù)處理（剔除不能正常打開的圖片，打亂數(shù)據(jù)集等）；然后對數(shù)據(jù)集進(jìn)行分割，其中90%的數(shù)據(jù)集作為train訓(xùn)練，另外10%的數(shù)據(jù)集作為test測試。

# 分割數(shù)據(jù)集，將全部數(shù)據(jù)分成0.9的Train和0.1的Test
source_path = r"./kagglecatsanddogs_5340/PetImages/"
# 如果不存在文件夾要新建一個(gè)
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目錄下移動(dòng)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"))
            
## 開始移動(dòng)，移動(dòng)前先剔除不能正常打開的圖片
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))
 
    #移動(dòng)10%文件到對應(yīng)目錄
    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:  # 只有未移動(dòng)過才需要移動(dòng)，否則每運(yùn)行一次都會(huì)移動(dòng)一下
        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、將數(shù)據(jù)轉(zhuǎn)成pytorch標(biāo)準(zhǔn)的DataLoader輸入格式

1、先對數(shù)據(jù)集進(jìn)行預(yù)處理，包括resize成224*224的尺寸，因?yàn)関gg_net模型需要的輸入尺寸為[N, 224, 224, 3]；隨機(jī)翻轉(zhuǎn)，隨機(jī)旋轉(zhuǎn)等，另外對數(shù)據(jù)集做Normalize標(biāo)準(zhǔn)化，其中的mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.2]是從ImageNet數(shù)據(jù)集上的百萬張圖片中隨機(jī)抽樣計(jì)算得到的，以上這些內(nèi)容主要是數(shù)據(jù)增強(qiáng)，增強(qiáng)模型的泛化性，有更好的預(yù)測效果。
2、然后將預(yù)處理好的數(shù)據(jù)轉(zhuǎn)成pytorch標(biāo)準(zhǔn)的DataLoader輸入格式，。

# 數(shù)據(jù)預(yù)處理
transform = transforms.Compose([
    transforms.RandomResizedCrop(224),# 對圖像進(jìn)行隨機(jī)的crop以后再resize成固定大小
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.2]),  # ImageNet全部圖片的平均值和標(biāo)準(zhǔn)差
    transforms.RandomRotation(20), # 隨機(jī)旋轉(zhuǎn)角度
    transforms.RandomHorizontalFlip(p=0.5), # 隨機(jī)水平翻轉(zhuǎn)
])
 
# 讀取數(shù)據(jù)
root = source_path
train_dataset = datasets.ImageFolder(root + '/train', transform)
test_dataset = datasets.ImageFolder(root + '/test', transform)
 
# 導(dǎo)入數(shù)據(jù)
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、導(dǎo)入預(yù)訓(xùn)練模型，并修改分類層

1、定義device，如果有GPU模型訓(xùn)練會(huì)自動(dòng)用GPU訓(xùn)練，否則會(huì)使用CPU；使用GPU訓(xùn)練，只需在模型、數(shù)據(jù)、損失函數(shù)上使用cuda()就行。
2、這邊默認(rèn)對分類圖像算法都熟悉，可以自己構(gòu)建vgg16的完整網(wǎng)絡(luò)，在貓狗數(shù)據(jù)集上重新訓(xùn)練。也可以下載預(yù)訓(xùn)練模型，由于原網(wǎng)絡(luò)的分類輸出是1000類別的，但是我們的圖片只有兩類，所以需要修改分類層，讓模型能夠適配我們的訓(xùn)練數(shù)據(jù)集。

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)  # 拿一個(gè)隨機(jī)tensor測試一下網(wǎng)絡(luò)的輸出是否滿足預(yù)期
output = vgg16(inputs.to(device))
print("原始VGG網(wǎng)絡(luò)的輸出：",output.size())

# 構(gòu)建新的全連接層
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("新構(gòu)建的VGG網(wǎng)絡(luò)的輸出：",output.size())

6、開始模型訓(xùn)練

開始模型訓(xùn)練，我們這里只訓(xùn)練全連接分類層，將特征層的梯度requires_grad設(shè)置為False，特征層的參數(shù)將不參與訓(xùn)練。
訓(xùn)練過程中保存效果最好的網(wǎng)絡(luò)模型，以防掉線，可以從斷點(diǎn)開始繼續(xù)訓(xùn)練，同時(shí)也可以用來做預(yù)測。
訓(xùn)練完成后，保存訓(xùn)練好的網(wǎng)絡(luò)和參數(shù)，后面可以加載模型做預(yù)測。

writer = SummaryWriter("./logs/model")
loss_func = nn.CrossEntropyLoss().to(device)
learning_rate = 0.0001

#如果我們想只訓(xùn)練模型的全連接層
for param in vgg16.features.parameters():
    param.requires_grad = False
optimizer = torch.optim.Adam(vgg16.parameters(),lr=learning_rate)

## 斷點(diǎn)續(xù)訓(xùn)開始
resume = False
if resume:
    # 恢復(fù)上次的訓(xùn)練狀態(tài)
    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

##訓(xùn)練開始
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)


    ## 測試開始，看訓(xùn)練效果是否滿足預(yù)期
    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))

    # 保存模型和訓(xùn)練參數(shù)的全相關(guān)信息，方便斷點(diǎn)續(xù)訓(xùn)
    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、利用訓(xùn)好的模型做預(yù)測

拿出一張圖片做預(yù)測，首先導(dǎo)入預(yù)訓(xùn)練模型，同樣改掉分類層，然后導(dǎo)入預(yù)訓(xùn)練權(quán)重，預(yù)測圖片類別，輸出標(biāo)簽值和預(yù)測類別。文章來源地址http://www.zghlxwxcb.cn/news/detail-423360.html

import matplotlib.pyplot as plt
img_path = r"./kagglecatsanddogs_5340/PetImages/test/Cat/1381.jpg"   # 拿出要預(yù)測的圖片
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("預(yù)測值為：",output)
print("預(yù)測標(biāo)簽為：",output.argmax(1).item())
print("預(yù)測動(dòng)物為：",train_dataset.classes[output.argmax(1)])

到了這里，關(guān)于Pytorch自定義數(shù)據(jù)集模型完整訓(xùn)練流程的文章就介紹完了。如果您還想了解更多內(nèi)容，請?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！