pytorch学习（一）——图像多分类实例

在这里我们使用一个图像的多分类来做例子，使我们对pytorch训练的流程进行一个简单的了解。
我使用的torch库的环境如下

torch==1.8.1+cu101
torchvision==0.9.1+cu101

1. 数据准备

这里我们采用牛津大学的102中花卉数据作为图像多分类的数据集，数据集在这里（提取码：1234），可以看到，数据分为 train 和 valid 两个文件夹，每个文件夹下面都有102个分类，每一个分类都有对应的图片。

2. 整体思路

在该项目中，我们应该先准备数据，一般图像训练都会进行数据增强，小批量的拿到数据，对图片数据设置一个管道，使图片数据能够自动从管道中流出，然后构建模型，这里我们使用预训练的残差神经网络模型来对图片进行训练，初始化一个优化器，构建一个损失函数，使模型在一次次迭代中能够不断进行优化，并存储最优的模型。该示例中导包如下所示：

import torch
from torchvision import transforms, models, datasets
from torch import nn, optim
import copy
import datetime

3. 数据增强

一般来说，图像数据为了防止过拟合以及增加能够训练的数据的量，都会使用数据增强，比如通过将原来的图像旋转，翻转，调节饱和度亮度等操作，将一张可训练的图片变为多张可训练的图片，在这里，我们对训练集和验证集分别进行变换，变换代码如下

def data_augmentation():
    '''
    数据增强
    :return: 含数据增强操作的变换器
    '''
    data_transform = {
        'train': transforms.Compose([
            transforms.RandomRotation(45),  # 随机旋转，角度在-45到45度之间
            transforms.CenterCrop(224),  # 从中心开始剪裁
            transforms.RandomHorizontalFlip(p=0.5),  # 以0.5的概率水平翻转
            transforms.RandomVerticalFlip(p=0.5),  # 以0.5的概率垂直翻转
            transforms.ColorJitter(brightness=0.2, contrast=0.1, saturation=0.1, hue=0.1),  # 参数依次为亮度、对比度、饱和度、色相
            transforms.RandomGrayscale(p=0.025),  # 以0.025的概率变为灰度图像，3通道即R=G=B
            transforms.ToTensor(),  # 将0-255的像素进行归一化
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 使用均值和标准差标准化三个通道的数据
        ]),
        'test': transforms.Compose([  # 测试集如果也进行归一化则参数必须和训练集一致
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    }
    return data_transform
    
transform = data_augmentation()

注意：验证集的变换必须和测试集保持一致，在这里的话表示为验证集的输入也必须是224大小，当然也可以直接 Resize(224)。

4. 构建管道

pytorch使用 Dataset 以及 DataLoader 来构建管道。

在这里先找到数据的训练集文件夹以及验证集文件夹，

data_dir = r'F:\机器学习\MyStudy\pytorch学习\data\102flowers\\'
train_dir = data_dir + 'train'
test_dir = data_dir + 'valid'

接下来就是使用这些数据来进行管道的构建了。

注意到，pytorch中有 ImageFolder 函数，该函数能够将传入路径的每一个文件夹当做一个分类，自动读取每个文件夹下的图片，这个函数正好能够拿过来用。

定义数据管道构建函数如下：

def data_load(train_dir, test_dir, batch_size, data_transform):
    '''
    构建数据管道
    :param train_dir: 训练集所在文件夹
    :param test_dir: 验证集所在文件夹
    :param batch_size: 每次迭代的批量大小
    :param data_transform: 数据增强转换器
    :return: 返回数据管道
    '''
    # 读取数据集
    image_datasets = {
        'train': datasets.ImageFolder(train_dir, data_transform['train']),
        'test': datasets.ImageFolder(test_dir, data_transform['test'])
    }
	# 构建管道
    dataloaders = {
        'train': torch.utils.data.DataLoader(image_datasets['train'], batch_size=batch_size, shuffle=True),
        'test': torch.utils.data.DataLoader(image_datasets['test'], batch_size=batch_size, shuffle=True)
    }
    return dataloaders

dataloaders = data_load(train_dir, test_dir, 16, transform)

5. 构建模型

模型方面我们使用预训练模型残差神经网络，该模型内置在 torchversion 的 models 模块中，但是，我们不可能原封不动的使用别人的预训练模型来进行使用，就比如最后一层，预训练模型是1000个输出，但是我们这里只需要102个类别的输出就行了，这里我们将最后一层输出层的类别数进行改变即可，模型的其他层我们将其冻结，即设置为参数不变。

def get_model(features):
    '''
    构建模型
    :param features: 最后需要分类的类别数
    :return: 模型和需要训练的参数
    '''
    # 默认使用残差神经网络的预训练模型
    model = models.resnet152(pretrained=True)

    for params in model.parameters():
        # 冻结模型每一层
        params.requires_grad = False

    # 改变最后一层线性层的输出大小
    # .fc等参数都是根据打印的模型来看出来的
    num_origin = model.fc.in_features
    model.fc = nn.Sequential(nn.Linear(num_origin, features), nn.LogSoftmax(dim=1))

	# 寻找可训练的参数，优化器需要传入可训练的参数
    param_learn = []
    for name, param in model.named_parameters():
        if param.requires_grad == True:
            param_learn.append(param)
    return model, param_learn

model, params_learn = get_model(102)

6. 设置优化器

优化器规定了需要更新的参数以及学习率等的值，选择的优化器对我们的训练结果影响十分大，一般的选择为 Adam 优化器，综合了动量梯度下降和RMSprop，除此之外，还可以设置动态学习率，每几个epoch学习率变为原来的几分之几等等。

# 设置优化器对哪些参数进行优化
optimizer = optim.Adam(params_learn, lr=1e-2)
# 学习率衰减，每7个epoch学习率衰减为原来的0.1倍
sch = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

7. 设置损失函数

损失函数的选取至关重要，这里我们选择softmax+交叉熵的形式，即最后一层是softmax，损失函数是交叉熵，代码中表示如下：

# 损失函数
criterion = nn.CrossEntropyLoss()

8. 训练并保存模型

所有的准备都做好后，就可以开始训练了。训练数据，就是从训练的数据管道中取出数据，进行迭代，反向传播，然后每训练一次使用测试集的管道进行一次验证，然后输出每次的准确率和损失就行。

def train_best(model, num_epoch, dataloaders, optimizer, loss_function, save_path, log_step=5):
    '''
    训练出最好的模型
    :param model: 构建的模型
    :param num_epoch: 总的要训练的轮次
    :param dataloaders: 数据管道
    :param optimizer: 优化器
    :param loss_function: 损失函数
    :param save_path: 保存的最优模型的路径
    :param log_step: 默认每5个训练batch打印一次数据
    :return: 每一个epoch的信息
    '''
    best_acc = 0
    # 最优模型
    best_model = copy.deepcopy(model.state_dict())

    # 保存每一个epoch的信息
    dfhistory = pd.DataFrame(columns=["epoch", "train_loss", "train_acc", "val_loss", "val_acc"])

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    print("Start Training...\n")
    nowtime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    print("==========" * 8 + "%s\n" % nowtime)

    for i in range(num_epoch):

        # 1，训练循环----------------------------------------------------------------

        loss_sum = 0.0
        metric_sum = 0.0
        # 训练模式，可以更新参数
        model.train()
        # 将数据全部取完

        # 记录每一个batch
        step = 0
        # 记录取了多少个数据
        all_step = 0

        for inputs, labels in dataloaders['train']:
            inputs = inputs.to(device)
            labels = labels.to(device)
            # 梯度清零，防止累加
            optimizer.zero_grad()

            # 每一批次拿了多少张图像
            a = inputs.size(0)

            outputs = model(inputs)
            loss = loss_function(outputs, labels)
            # 返回每一行的最大值和其索引
            _, pred = torch.max(outputs, 1)
            loss.backward()
            optimizer.step()

            # 学习损失
            loss_sum += loss.item() * inputs.size(0)
            metric_sum += torch.sum(pred == labels.data)

            step += 1
            all_step += a

            if step % log_step == 0:
                print("[step = {}]  train_loss = {:.3f}, train_acc = {:.3f}".
                      format(all_step, loss_sum / all_step, metric_sum.double() / all_step))

        train_loss = loss_sum / len(dataloaders['train'].dataset)
        train_acc= metric_sum.double() / len(dataloaders['train'].dataset)

        # 2，验证循环----------------------------------------------------------------

        val_loss_sum = 0.0
        val_metric_sum = 0.0

        step = 0
        all_step = 0

        # 验证模式，该模式下模型参数不能进行修改
        model.eval()

        for inputs, labels in dataloaders['test']:
            inputs = inputs.to(device)
            labels = labels.to(device)
            # 梯度清零，防止累加
            optimizer.zero_grad()

            a = inputs.size(0)

            outputs = model(inputs)
            loss = loss_function(outputs, labels)
            # 返回每一行的最大值和其索引
            _, pred = torch.max(outputs, 1)

            # 学习损失
            val_loss_sum += loss.item() * inputs.size(0)
            val_metric_sum += torch.sum(pred == labels.data)

            step += 1
            all_step += a

            if step % log_step == 0:
                print("[step = {}]  val_loss = {:.3f}, val_acc = {:.3f}".
                      format(all_step, val_loss_sum / all_step, val_metric_sum.double() / all_step))

        val_loss = val_loss_sum / len(dataloaders['test'].dataset)
        val_acc = val_metric_sum.double() / len(dataloaders['test'].dataset)

        # 3. 打印epoch级别日志
        print("EPOCH = {}/{}  train_loss = {:.3f}, train_acc = {:.3f}, val_loss = {:.3f}, val_acc = {:.3f}\n".
              format(i, num_epoch, train_loss, train_acc, val_loss, val_acc))

        nowtime = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print("==========" * 8 + "%s\n" % nowtime)

        # 4. 保存最优的模型
        if val_acc > best_acc:
            best_acc = val_acc
            best_model = copy.deepcopy(model.state_dict())
            state = {
                'state_dict': model.state_dict(),
                'best_acc': best_acc,
                'optimizer': optimizer.state_dict()
            }
            torch.save(state, save_path)
		
		# 记录每个epoch的信息
        dfhistory.loc[i] = (i, train_loss, train_acc, val_loss, val_acc)
        
    print('Finished Training...\n')
    
    return dfhistory

进行训练

train_best(model, 5, dataloaders, optimizer, criterion, 'model.pth')

好了，经过上面的图像多分类过程，相信大家对pytorch进行深度学习已经有了一个大概的了解了，而且以后进行类似的图像多分类任务也可以直接套用上面的函数，那么下面我们会对每一步进行一些详细的讲解。