文章目录

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 torchimport torchvisionfrom torchvision import datasets, transformsimport torch.utils.dataimport matplotlib.pyplot as pltfrom torch.utils.data import TensorDataset,DataLoader,Datasetfrom torch.utils.tensorboard import SummaryWriterimport torch.nn.functional as Ffrom torch import nnimport numpy as npimport osimport shutilfrom PIL import Imageimport warningswarnings.filterwarnings("ignore")

3、分割数据集

下载猫狗大战数据集，并解压。
解压完成后，通过以下代码实现数据集预处理（剔除不能正常打开的图片，打乱数据集等）；然后对数据集进行分割，其中90%的数据集作为train训练，另外10%的数据集作为test测试。

# 分割数据集，将全部数据分成0.9的Train和0.1的Testsource_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_pathtrain_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 = Falseoptimizer = 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 = 0total_test_step = 0min_acc = 0.0for 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 pltimg_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)])