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

2、导入各种需要用到的包

import torch   //用于导入名为"torch"的模块。torch 是一个广泛使用的库，用于构建和训练神经网络。它提供了丰富的功能和工具，包括张量操作、自动求导、优化算法等，使得深度学习任务更加简单和高效。可以使用torch.Tensor类来创建张量，使用torch.nn.Module类来定义神经网络模型，使用torch.optim模块来选择优化算法等等。
import torchvision //通过导入 torchvision 模块，你可以使用其中的类、函数和数据集来处理图像数据、加载预训练模型、进行数据增强等等。例如，你可以使用 torchvision.transforms 模块来定义图像的转换操作，使用 torchvision.models 模块来加载预训练的神经网络模型，使用 torchvision.datasets 模块来加载标准的计算机视觉数据集等等
from torchvision import datasets, transforms//通过从torchvision模块中导入datasets和transforms类实现了对这两个类的引用。

• datasets类提供了许多常见的计算机视觉数据集，如MNIST、CIFAR-10等。你可以使用datasets类加载这些数据集，并在深度学习模型中使用它们进行训练和评估。

• transforms类提供了一系列的图像变换操作，用于在加载图像数据时进行预处理。你可以使用transforms类中的方法对图像进行常见的变换，如缩放、裁剪、翻转、归一化等。

import torch.utils.data//torch.utils.data 模块是PyTorch中用于处理数据加载和预处理的工具包。它提供了一系列的类和函数，用于构建数据加载器（DataLoader），定义数据集（Dataset），以及进行数据的批处理、采样、转换等操作。
import matplotlib.pyplot as plt//matplotlib.pyplot 是一个用于绘制图表和可视化数据的模块。它是Matplotlib库的一个子模块，提供了简单而强大的绘图功能，可以绘制各种类型的图形，如折线图、散点图、柱状图、饼图等。
from torch.utils.data import TensorDataset,DataLoader,Dataset//从torch.utils.data模块中导入三个类：TensorDataset、DataLoader和Dataset。

• TensorDataset是一个用于包装张量数据的数据集类。它将一个或多个张量作为输入，并将它们打包在一起，可以用于在训练过程中提供数据和标签。

• • DataLoader是一个数据加载器类，用于批处理和并行加载数据。它可以从一个数据集（如TensorDataset）中加载数据，并按照指定的批次大小进行批处理。DataLoader还可用于设置数据加载的参数，如并行加载、数据打乱等。Dataset是一个抽象基类，用于自定义数据集。如果你想创建自己的数据集类，可以继承Dataset类并实现__len__和__getitem__方法，以定义数据集的长度和访问方式。

from torch.utils.tensorboard import SummaryWriter//SummaryWriter是PyTorch中用于将训练过程中的数据传输到TensorBoard进行可视化的类。TensorBoard是一个用于可视化和监控深度学习模型的工具，可以显示训练过程中的损失曲线、参数分布、图像、计算图等。
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、分割数据集

# 分割数据集，将全部数据分成0.9的Train和0.1的Test
source_path = r"./kagglecatsanddogs_5340/PetImages/" //定义了一个名为source_path的变量，并给它赋值为"./kagglecatsanddogs_5340/PetImages/"。路径字符串前面带有r前缀，表示这是一个原始字符串（raw string），这意味着字符串中的反斜杠\将被视为普通字符，而不是转义字符。这在表示文件路径时非常方便，因为在Windows系统中，文件路径常常包含反斜杠。
# 如果不存在文件夹要新建一个
if not os.path.exists(os.path.join(source_path, "train")):  //

os.path.join(source_path, "train")代码将"train"加入到source_path路径中，可以得到表示"train"文件夹的完整路径。

然后，os.path.exists()函数用于检查给定路径是否存在。如果路径存在，返回True；如果路径不存在，返回False。

os.mkdir(os.path.join(source_path, "train")) //

将source_path和"train"进行路径拼接，得到的是指向"train"文件夹的完整路径。

然后，os.mkdir()函数将这个完整路径作为参数传递进去，创建了一个名为"train"的文件夹。

train_dir = os.path.join(source_path, "train")  //os.path.join()函数将source_path和"train"组合起来形成一个新的路径，并将其赋值给变量train_dir。

if not os.path.exists(os.path.join(source_path, "test")):   //使用了os.path模块来检查指定路径下的"test"文件夹是否存在。
os.mkdir(os.path.join(source_path, "test"))  //这段代码使用了os.mkdir()函数来创建一个名为"test"的文件夹，路径是由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)])