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
1
2
3
4
5

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()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

注意：验证集的变换必须和测试集保持一致，在这里的话表示为验证集的输入也必须是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'
1
2
3

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

注意到，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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

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)
1
2
3
4

7. 设置损失函数

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

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

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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

进行训练

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

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

全部的代码可以在GitHub仓库进行查看。