使用pytorch训练自己的数据集

一、introduction

不得不说pytorch的功能十分强大，其中torchvision.datasets已经内置了常用的数据集，我们也可以使用datasets.ImageFolder来加载我们自己的数据集。如果你觉得这两个方法都不能很好的加载你想要的数据集，也可以通过继承dataset类来自定义自己的数据集加载方式。

本文主要介绍前两种加载数据集的方法，然后使用常见的卷积神经网络来训练并测试自己的数据集。

二、加载数据集

2.1 加载FashionMNIST数据集

def load_FashionMNIST(root,batch_size):
    mnist_train = torchvision.datasets.FashionMNIST(root=root, train=True, download=True,transform=transforms.ToTensor())
    mnist_test = torchvision.datasets.FashionMNIST(root=root, train=False, download=True,transform=transforms.ToTensor())
    
    train_iter = torch.utils.data.DataLoader(mnist_train, batch_size=batch_size, shuffle=True)
    test_iter = torch.utils.data.DataLoader(mnist_test, batch_size=batch_size, shuffle=False)

    return train_iter,test_iter

通过代码可以看到加载FashionMNIST数据集非常简单，只需要通过一行代码即可。其中torchvision.datasets.FashionMNIST(）中参数的含义如下：
root表示存放本地存放MNIST数据集的文件夹地址。
train如果为True，则使用训练数据集，否则使用测试数据集。
download如果为true，则从网上下载数据集并将其放在根目录root中。如果数据集已下载，则不会再次下载。
transform是图像转换的方法，能够实现数据增强并将其转换成pytorch中形状为(C,H, W)的Tensor格式，且/255归一化到[0,1.0]之间。（转换方法需要自己定义）

获取到数据集之后通过 torch.utils.data.DataLoader方法获取到每个batch的迭代器。该方法中第一个参数为原始数据，类型为 dataset，第二个参数batch_size为批处理大小，第三个参数shuffle决定是否打乱数据集。

2.2 加载自己的数据集

这里我们使用torchvision.datasets.ImageFolder方法来加载自己的数据集。这是pytorch的一个通用的数据集加载器，在加载数据集时要注意的是本地数据集中图片的存放方式如下图所示（参考地址），其中文件夹名称就是其对应的标签。

对于数据集中训练集和测试集分开和没有分开的两种情况，分别使用以下代码：

def load_local_dataset(dataset_dir, ratio = 0.8, batch_size = 256):
    #获取数据集
    all_datasets = datasets.ImageFolder(dataset_dir, transform=train_transforms)
    #将数据集划分成训练集和测试集
    train_size=int(ratio * len(all_datasets))
    test_size=len(all_datasets) - train_size
    train_datasets, test_datasets = torch.utils.data.random_split(all_datasets, [train_size, test_size])

    train_iter = torch.utils.data.DataLoader(train_datasets, batch_size=batch_size, shuffle=True)
    test_iter = torch.utils.data.DataLoader(test_datasets, batch_size=batch_size, shuffle=True)

    return train_iter,test_iter

def load_train_test_dataset(train_dir, test_dir , batch_size = 256):
    #获取数据集
    train_datasets = datasets.ImageFolder(train_dir, transform=train_transforms)
    test_datasets = datasets.ImageFolder(test_dir, transform=test_transforms)

    train_iter = torch.utils.data.DataLoader(train_datasets, batch_size=batch_size, shuffle=True)
    test_iter = torch.utils.data.DataLoader(test_datasets, batch_size=batch_size, shuffle=True)

    return train_iter,test_iter

最后补上图像转换（transform）方法和头文件：

import torch
import torchvision
from torchvision import datasets, transforms

#对训练集做一个变换
train_transforms = transforms.Compose([
    transforms.RandomResizedCrop(224),		#对图片尺寸做一个缩放切割
    transforms.RandomHorizontalFlip(),		#水平翻转
    transforms.ToTensor(),					#转化为张量
    transforms.Normalize((.5, .5, .5), (.5, .5, .5))	#进行归一化
])

#对测试集做变换
test_transforms = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.ToTensor(),
    transforms.Normalize((.5, .5, .5), (.5, .5, .5))
])

三、训练和测试

首先加载模型、数据集，定义损失函数和优化函数。在使用代码训练自己的数据集时记得修改地址dataset_dir 。

import torch
import load_dataset

#加载模型
import ResNet
net = ResNet.ResNet_model(50).cuda()  #使用GPU训练
#net = torch.nn.DataParallel(ResNet.net).cuda()  #使用多块GPU共同训练
print("model loaded...")

#加载数据集
ratio=0.8
batch_size=512

#加载MNIST数据集，图片大小为28x28x1，记得修改网络结构
# root="E:/数据集"
# train_iter,test_iter=load_dataset.load_FashionMNIST(root,batch_size)

#加载训练集和测试集
# train_dir = "E:/数据集/rice_diseases/train"
# test_dir = "E:/数据集/rice_diseases/test"
# train_iter, test_iter = load_dataset.load_train_test_dataset(train_dir, test_dir, batch_size)

#训练集和测试集在一个文件夹下
dataset_dir = "../dataset/kitchen_waste/train"
train_iter,test_iter=load_dataset.load_local_dataset(dataset_dir,ratio,batch_size)

print("data loaded...")
print("训练集=",len(train_iter))
print("测试集=",len(test_iter))

#定义损失函数和优化器
lr,num_epochs =0.001, 1000
loss = torch.nn.CrossEntropyLoss().cuda()  #损失函数
optimizer = torch.optim.Adam(net.parameters(), lr=lr) #优化器

接下来进行训练和测试：

#训练模型
def train(net, train_iter, test_iter, optimizer,  loss, num_epochs):
    for epoch in range(num_epochs):
        # 训练过程
        net.train()  # 启用 BatchNormalization 和 Dropout
        train_l_sum, train_acc_sum, train_num = 0.0, 0.0, 0
        for X, y in train_iter:
            X = X.cuda()
            y = y.cuda()
            y_hat = net(X)
            l = loss(y_hat, y).sum()
            optimizer.zero_grad()
            l.backward()
            optimizer.step()
            #计算准确率
            train_l_sum += l.item()
            train_acc_sum += (y_hat.argmax(dim=1) == y).sum().item()
            train_num += y.shape[0]
        print('epoch %d, loss %.4f, train acc %.3f' % (epoch + 1, train_l_sum / train_num, train_acc_sum / train_num))

        # 测试过程
        if (epoch+1) %5 == 0:
            test_acc_sum, test_num= 0.0, 0
            with torch.no_grad(): #不会求梯度、反向传播
                net.eval()  # 不启用 BatchNormalization 和 Dropout
                for X,y in test_iter:
                    X = X.cuda()
                    y = y.cuda()
                    test_acc_sum += (net(X).argmax(dim=1) == y).float().sum().item()
                    test_num += y.shape[0]
                print('test acc %.3f' % (test_acc_sum / test_num))

            torch.save(net.module.state_dict(), f'./checkpoint/model_{str(epoch + 1).zfill(4)}.pt')  # 保存模型

train(net, train_iter, test_iter, optimizer, loss, num_epochs)

这里默认使用了GPU进行训练，如果你想用CPU进行训练，把所有.cuda（）后缀删除即可。在进行测试时默认每隔五代测试一次，并保存模型参数。

四、writing in the end

以上所有代码均经过博主运行测试，没有任何问题。如果你在运行代码的过程中有任何疑问，或者关于代码有不理解的地方，均可以在底下评论区留言。

在以上代码中有加载模型的代码，关于网络结构的详细内容请看：
使用pytorch搭建自己的网络之AlexNet
使用pytorch搭建自己的网络之VGG
使用pytorch搭建自己的网络之GoogLeNet
使用pytorch搭建自己的网络之ResNet