跳转至

计算机视觉迁移学习

在本教程中，我们将介绍迁移学习的基本原理，并展示一个在计算机视觉领域中的迁移学习的使用示例。

原理简介

迁移学习 (Transfer Learning) 是一种将从源数据集学到的知识迁移到目标数据集的方法。

众所周知，有监督学习是一种相当常见的深度学习模型的训练方式，但它需要大量带标注的数据才能达到较好的效果，当我们想将某个模型应用于某个特定的任务上时，通常受制于成本而无法获得大量带标注的数据，如果直接在这样的小规模数据上进行训练，很容易造成过拟合。而迁移学习是解决这一问题的方法之一。

以计算机视觉领域中常见的图像分类任务为例，一般的图像分类模型可以分为两个部分：特征提取器（或称为骨干网络）和分类器（或称为输出层）。特征提取器一般是诸如卷积神经网络的多层网络，分类器一般是诸如全连接层的单层网络。由于不同分类任务的类别一般不同，分类器通常无法复用，而特征提取器通常可以复用，虽然源数据集中的物体可能与目标数据集大相径庭，甚至完全没有交集，但在大规模数据上预训练得到的模型可能具备提取更常规的图像特征（例如边缘、形状和纹理）的能力，从而有助于有效地识别目标数据集中的物体。

假设我们已有一个预训练模型，大致有三种使用方式：

  1. 使用预训练模型的参数对特征提取器进行初始化，然后对整个模型进行训练。 对于深度学习模型来说，参数初始化的方法对保持数值稳定性相当重要，不当的初始化方法可能会导致在训练时出现梯度爆炸或梯度消失的问题。如果使用预训练模型进行初始化，可以在很大程度上保证模型参数初始值的合理性，让模型“赢在起跑线上”。

  2. 对整个模型进行训练，但对特征提取器使用较小的学习率，对分类器使用较大的学习率。 预训练得到的特征提取器已经得到了充分的训练，所以只需要较小的学习率；而分类器的参数通常是随机初始化的，所以需要从头开始学习，因此需要较大的学习率。

  3. 固定特征提取器的参数，只训练分类器。 如果目标数据集的类别恰好是源数据集的子集，那么这样的方式一般会很有效且快速。

迁移学习示例

在本节中，我们将使用 ResNet-18 作为特征提取器在 CIFAR-10 数据集 上进行图像分类任务。

ResNet-18 的预训练模型（在 ImageNet 数据集上训练得到）和 CIFAR-10 数据集都可以通过 FlowVision 方便地获取。

首先导入所需的依赖：

import oneflow as flow
from oneflow import nn
from oneflow.utils.data import DataLoader

from flowvision.models import resnet18
from flowvision.datasets import CIFAR10
import flowvision.transforms as transforms

定义 epoch, batch size, 以及使用的计算设备： 

NUM_EPOCHS = 3
BATCH_SIZE = 64
DEVICE = 'cuda' if flow.cuda.is_available() else 'cpu'

数据加载及预处理

定义 Dataset 和 DataLoader:

train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

train_dataset = CIFAR10(root='./data', train=True, transform=train_transform, download=True)
test_dataset = CIFAR10(root='./data', train=False, transform=test_transform, download=True)

train_data_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
test_data_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)

定义模型

model = resnet18(pretrained=True)
在这里，我们通过将 pretrained 参数设置为 True 来获取加载了预训练权重的 ResNet-18 模型。如果输出 model.fc，将会得到 "Linear(in_features=512, out_features=1000, bias=True)"，可以看出此分类器有 1000 个输出神经元，对应于 ImageNet 的 1000 个类别。CIFAR-10 数据集是 10 个类别，因此我们需要替换掉这个全连接层分类器：

model.fc = nn.Linear(model.fc.in_features, 10)

然后将模型加载到计算设备： 

model = model.to(DEVICE)

训练模型

定义训练函数： 

def train_model(model, train_data_loader, test_data_loader, loss_func, optimizer):
    dataset_size = len(train_data_loader.dataset)
    model.train()
    for epoch in range(NUM_EPOCHS):
        for batch, (images, labels) in enumerate(train_data_loader):
            images, labels = images.to(DEVICE), labels.to(DEVICE)
            preds = model(images)
            loss = loss_func(preds, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if batch % 100 == 0:
                print(f'loss: {loss:>7f}  [epoch: {epoch} {batch * BATCH_SIZE:>5d}/{dataset_size:>5d}]')

        evaluate(model, test_data_loader)

定义评估函数，使用准确率作为评估指标： 

def evaluate(model, data_loader):
    dataset_size = len(data_loader.dataset)
    model.eval()
    num_corrects = 0
    for images, labels in data_loader:
        images, labels = images.to(DEVICE), labels.to(DEVICE)
        preds = model(images)
        num_corrects += flow.sum(flow.argmax(preds, dim=1) == labels)

    print('Accuracy: ', num_corrects.item() / dataset_size)

我们可以通过给优化器传入相应的需要优化的参数，来实现上文中提到的三种方式。

第 1 种方式，对整个模型进行训练：

optimizer = flow.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4)

第 2 种方式，对特征提取器使用较小的学习率，对分类器使用较大的学习率：

fc_params = list(map(id, model.fc.parameters()))
backbone_params = filter(lambda p: id(p) not in fc_params, model.parameters())
optimizer = flow.optim.SGD([{'params': backbone_params, 'lr': 0.0001},
                            {'params': model.fc.parameters(), 'lr': 0.001}],
                            momentum=0.9, weight_decay=5e-4)

第 3 种方式，固定特征提取器的参数，只训练分类器：

optimizer = flow.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9, weight_decay=5e-4)

开始训练：

loss_func = nn.CrossEntropyLoss()
train_model(model, train_data_loader, test_data_loader, loss_func, optimizer)

结果对比

在使用迁移学习的情况下（这里使用第一种方式），模型在经过 3 个 epoch 的训练后在测试集上的准确率达到了 0.9017; 如果从头开始训练、不使用迁移学习，同样经过 3 个 epoch 的训练后，准确率仅为 0.4957。这表明迁移学习确实能起到显著的作用。

Back to top