当前位置: 首页 > news >正文

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

文章目录

  • Pytorch模型自定义数据集训练流程
    • 1、任务描述
    • 2、导入各种需要用到的包
    • 3、分割数据集
    • 4、将数据转成pytorch标准的DataLoader输入格式
    • 5、导入预训练模型,并修改分类层
    • 6、开始模型训练
    • 7、利用训好的模型做预测

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

我们以kaggle竞赛中的猫狗大战数据集为例搭建Pytorch自定义数据集模型训练的完整流程。

1、任务描述

Cats vs. Dogs(猫狗大战)数据集是Kaggle大数据竞赛某一年的一道赛题,利用给定的数据集,用算法实现猫和狗的识别。 其中包含了训练集和测试集,训练集中猫和狗的图片数量都是12500张且按顺序排序,测试集中猫和狗混合乱序图片一共12500张。

在这里插入图片描述

卷积神经网络(CNN)是一类包含卷积计算且具有深度结构的前馈神经网络,是深度学习的代表算法之一。卷积神经网络具有表征学习能力,能够按其阶层结构对输入信息进行平移不变分类,因此也被称为“平移不变人工神经网络”。不使用深度学习框架,用numpy基础代码来构建自己的深度学习网络。

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、分割数据集

下载猫狗大战数据集,并解压。
下载地址:https://www.kaggle.com/c/dogs-vs-cats/data。
解压完成后,通过以下代码实现数据集预处理(剔除不能正常打开的图片,打乱数据集);然后对数据集进行分割,其中0.9的数据集作为train训练,0.1的数据集作为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)

##训练开始
total_train_step = 0
total_test_step = 0
min_acc = 100.0
for epoch in range(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
    total_acc = 0.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)
    
    ## 保存Acc最小的模型
    if val_acc < min_acc:
        min_acc = val_acc
        torch.save(vgg16.state_dict(), "./models/2classes_vgg16_weight_{}_{}.pth".format(epoch, round(val_acc,4)))
        torch.save(vgg16, "./models/2classes_vgg16_{}_{}.pth".format(epoch, round(val_acc,4)))

    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)])

相关文章:

  • wordpress 本机安装目录/免费引流推广怎么做
  • 什么做网站做个网站一般要多少钱啊/网站seo优化服务
  • 福州做网站哪家好/日本网站源码
  • 网站图片翻页效果如何做/中国国际新闻
  • 天津网站建设哪里好/网站域名查询ip地址
  • 北京市国家住房建设委员会网站/江苏疫情最新消息
  • 云原生技术学习笔记(基础版)
  • nohup + 命令实现后台不挂断地运行程序
  • 贪心算法专题
  • Android项目Gadle统一依赖管理
  • 水声功率放大器模块在圆柱壳结构声源辐射研究中的应用
  • uefi和legacy的区别对比
  • windows安装VMware最新版本(VMware Workstation 17.0 Pro)详细教程
  • 亚马逊云科技助力游戏上云学习心得-运行篇
  • LeetCode 334. 递增的三元子序列(C++)
  • VTK-vtkSelectPolyDataFilter
  • CSS选择器整理学习(上)
  • 我收集的PDF电子书