PyTorch项目笔记（三）使用ImageNet预训练ResNet18模型训练图像分类模型

1 加载ImageNet预训练模型

在torchvision.model包中定义了许多模型用于完成图像方面的深度学习任务，包括：

• 图像分类
• 语义分割（semantic segmentation）
• 目标检测（object detection）
• 实例分割（instance segmentation）
• 人物关键点检测（person keypoint detection）
• 视频分类

下面将给出部分预训练模型的例子：

import torchvision.models as models

resnet18 = models.resnet18()
alexnet = models.alexnet()
vgg16 = models.vgg16()
squeezenet = models.squeezenet1_0()
densenet = models.densenet161()
inception = models.inception_v3()
googlenet = models.googlenet()
shufflenet = models.shufflenet_v2_x1_0()
mobilenet = models.mobilenet_v2()
resnext50_32x4d = models.resnext50_32x4d()
wide_resnet50_2 = models.wide_resnet50_2()
mnasnet = models.mnasnet1_0()

当传入pretrained=True参数时，模型会从网络上下载状态字典并保存到本地磁盘，最终加载状态字典得到预训练的模型。

resnet18 = models.resnet18(pretrained=True)
print(resnet18)

使用print()函数可查看ResNet18模型的内部结构，结果如下：

ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential((0): BasicBlock(...) (1): BasicBlock(...))
  (layer2): Sequential((0): BasicBlock(...) (1): BasicBlock(...))
  (layer3): Sequential((0): BasicBlock(...) (1): BasicBlock(...))
  (layer4): Sequential((0): BasicBlock(...) (1): BasicBlock(...))
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  (fc): Linear(in_features=512, out_features=1000, bias=True)
)

此处省略BasicBlock内部构成，如感兴趣可自行输出进行观察。

2 准备数据集

本文使用小型数据集hymenoptera进行训练，该数据集是ImageNet的子集，只包含蚂蚁和蜜蜂两个类别。
需要下载数据集可点击链接：hymenoptera数据集

2.1 加载数据集

加载数据集代码如下所示：

from torchvision import datasets, transforms

# 在训练集上：扩充、归一化
# 在验证集上：归一化
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}
#数据集存放路径
data_dir = 'data/hymenoptera_data'
#使用torchvision.datasets.ImageFolder类快速封装数据集
#此处使用了lambda语法
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                          data_transforms[x])
                  for x in ['train', 'val']}
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
                                             shuffle=True, num_workers=4)
              for x in ['train', 'val']}
#读取数据集的数目
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
#读取数据集中的图像种类
class_names = image_datasets['train'].classes

2.2 使用matplotlib可视化数据集

由于一个batch中的图像是保存在tensor中，一张维度是[H,W,C]、值范围是[0,255]的图片需要经过ToTensor转换成维度是[C,H,W]、值范围是[0,1]的Tensor，在经过Normalize完成归一化。而使用matplotlib显示的图像需要对图像进行反向操作才能正常显示。

import matplotlib.pyplot as plt

def imshow(inp, title=None): 
    # 可视化一组 Tensor 的图片
    inp = inp.numpy().transpose((1, 2, 0)) #转换图片维度
    mean = np.array([0.485, 0.456, 0.406]) 
    std = np.array([0.229, 0.224, 0.225]) 
    inp = std * inp + mean #反向操作
    inp = np.clip(inp, 0, 1) 
    plt.imshow(inp) 
    if title is not None: 
        plt.title(title) 
    plt.pause(0.001) # 暂停一会儿，为了将图片显示出来
# 获取一批训练数据
inputs, classes = next(iter(dataloaders['train'])) 
# 批量制作网格
out = torchvision.utils.make_grid(inputs) 
imshow(out, title=[class_names[x] for x in classes]) 

运行结果如下图所示。

3 模型训练函数

def train_model(model, criterion, optimizer, scheduler, num_epochs=25): 
    """ 训练模型，并返回在验证集上的最佳模型和准确率 
    Args: 
    - model(nn.Module): 要训练的模型 
    - criterion: 损失函数 
    - optimizer(optim.Optimizer): 优化器 
    - scheduler: 学习率调度器 
    - num_epochs(int): 最大 epoch 数 
    Return: 
    - model(nn.Module): 最佳模型 
    - best_acc(float): 最佳准确率 
    """
    since = time.time()

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print(f'Epoch {epoch}/{num_epochs - 1}')
        print('-' * 10)

        # 训练集和验证集交替进行前向传播
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # 设置为训练模式，可以更新网络参数
            else:
                model.eval()   # 设置为预估模式，不可更新网络参数

            running_loss = 0.0
            running_corrects = 0

            # 遍历数据集
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # 清空梯度，避免累加了上一次的梯度
                optimizer.zero_grad()

                with torch.set_grad_enabled(phase == 'train'):
                    # 正向传播
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)

                    # 反向传播且仅在训练阶段进行优化
                    if phase == 'train':
                        loss.backward() # 反向传播
                        optimizer.step()

                # 统计loss、准确率
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)
            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]

            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')

            # 发现了更优的模型，记录起来
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print()

    time_elapsed = time.time() - since
    print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
    print(f'Best val Acc: {best_acc:4f}')

    # 加载训练的最好的模型
    model.load_state_dict(best_model_wts)
    return model

4 使用torchvision微调模型

微调是值在创建模型时，使用预训练模型来初始化网络，而非随机初始化网络，并且权重可以随着训练的进行而发生改变。对预训练模型进行微调的三个步骤：

1. 替换输出层。将ResNet18的最后一个全连接层（输出1000维）替换维新的全连接层（输出2维）。
2. 训练输出层。新的输出层会将前面的层提取出的低级特征映射待我们所期望的类别的概率。
3. 训练其余层。将这些层的权重标记为需要求导。
   具体代码如下所示：

model = models.resnet18(pretrained=True) # 加载预训练模型
num_ftrs = model.fc.in_features # 获取低级特征维度 
model.fc = nn.Linear(num_ftrs, 2) # 替换新的输出层 
model = model.to(device) 
# 交叉熵作为损失函数 
criterion = nn.CrossEntropyLoss() 
# 所有参数都参加训练 
optimizer_ft = optim.SGD(model.parameters(), lr=0.001, momentum=0.9) 
# 每过 7 个 epoch 将学习率变为原来的 0.1 
scheduler = optim.lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
model_conv = train_model(model_conv, criterion, optimizer_conv,
                         exp_lr_scheduler, num_epochs=5)

此处训练五轮为例，训练结果如下图所示。

5 观察模型预测结果

def visualize_model(model, num_images=6):
    was_training = model.training
    model.eval()
    images_so_far = 0
    fig = plt.figure()

    with torch.no_grad():
        for i, (inputs, labels) in enumerate(dataloaders['val']):
            inputs = inputs.to(device)
            labels = labels.to(device)

            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)

            for j in range(inputs.size()[0]):
                images_so_far += 1
                ax = plt.subplot(num_images//2, 2, images_so_far)
                ax.axis('off')
                ax.set_title(f'predicted: {class_names[preds[j]]}')
                imshow(inputs.cpu().data[j])

                if images_so_far == num_images:
                    model.train(mode=was_training)
                    return
        model.train(mode=was_training)
visualize_model(model_conv)

plt.ioff()
plt.show()

6 固定模型参数

微调预训练模型，需要修改模型的内部结构，使其符合具体任务。模型所用框架不一样，在 将其他框架编写的模型迁移到 PyTorch 中时，无法使它们兼容。此时可以采取 Pipeline 形式将预训 练模型的参数固定，或者说将前一个模型的输出保存下来，将该输出作为 PyTorch 模型的输入。

采取这种思路，我们可以将模型除了输出层之外的所有层看成一个特征提取器。在训练模型 时，这些层的权重不参与训练，不可优化。在 PyTorch 中将权重设置为不可训练，只需将 requires_grad 设置为 False 即可。例如，下述代码可以将 ResNet18 的所有层设置为不可训练。

model_conv = torchvision.models.resnet18(pretrained=True) # 加载预训练模型
for param in model_conv.parameters(): # 锁定模型所有参数
    param.requires_grad = False

num_ftrs = model_conv.fc.in_features  # 获取低级特征维度
model_conv.fc = nn.Linear(num_ftrs, 2) # 替换新的输出层

model_conv = model_conv.to(device)

criterion = nn.CrossEntropyLoss()

# 只有最后一层全连接层fc，参加训练
optimizer_conv = optim.SGD(model_conv.fc.parameters(), lr=0.001, momentum=0.9)

# 每过 7 个 epoch 将学习率变为原来的 0.1 
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)
model_conv = train_model(model_conv, criterion, optimizer_conv,
                         exp_lr_scheduler, num_epochs=5)
visualize_model(model_conv)

plt.ioff()
plt.show()