评价此页

★ ★ ★ ★ ★

beginner/transfer_learning_tutorial

在 Google Colab 中运行

注意

跳转至页面底部下载完整示例代码。

计算机视觉迁移学习教程#

创建日期：2017年3月24日 | 最后更新：2025年1月27日 | 最后验证：2024年11月5日

作者: Sasank Chilamkurthy

在本教程中，你将学习如何利用迁移学习训练一个用于图像分类的卷积神经网络。你可以在 cs231n 笔记中阅读更多关于迁移学习的内容。

引用其中的内容：

在实践中，很少有人从头开始（使用随机初始化）训练整个卷积神经网络，因为拥有足够大规模的数据集相对较少。相反，通常的做法是在一个非常大的数据集（例如 ImageNet，包含 120 万张图片和 1000 个类别）上预训练一个卷积网络，然后将该卷积网络用作初始化，或作为目标任务的固定特征提取器。

这两种主要的迁移学习场景如下：

微调卷积网络 (Finetuning the ConvNet)：我们不使用随机初始化，而是使用预训练网络来初始化网络（例如在 ImageNet 1000 数据集上训练好的网络）。其余的训练过程与往常一样。
卷积网络作为固定特征提取器 (ConvNet as fixed feature extractor)：在这里，我们将冻结网络中除最后一层全连接层以外所有层的权重。最后一层全连接层将被替换为一个新的、带有随机权重的层，并且只训练这一层。

# License: BSD
# Author: Sasank Chilamkurthy

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.backends.cudnn as cudnn
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
from PIL import Image
from tempfile import TemporaryDirectory

cudnn.benchmark = True
plt.ion()   # interactive mode

<contextlib.ExitStack object at 0x7f64b05c3370>

加载数据#

我们将使用 torchvision 和 torch.utils.data 包来加载数据。

我们要解决的问题是训练一个模型来分类蚂蚁和蜜蜂。我们大约有 120 张蚂蚁图片和 120 张蜜蜂图片作为训练集。每类有 75 张验证图片。通常情况下，如果从头开始训练，这个数据集太小，无法实现良好的泛化。由于我们使用了迁移学习，我们应该能够获得相当不错的泛化效果。

该数据集是 ImageNet 的一个非常小的子集。

注意

从此处下载数据并将其解压到当前目录。

# Data augmentation and normalization for training
# Just normalization for validation
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'
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

# We want to be able to train our model on an `accelerator <https://pytorch.ac.cn/docs/stable/torch.html#accelerators>`__
# such as CUDA, MPS, MTIA, or XPU. If the current accelerator is available, we will use it. Otherwise, we use the CPU.

device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
print(f"Using {device} device")

Using cuda device

可视化部分图像#

让我们可视化一些训练图像，以便理解数据增强。

def imshow(inp, title=None):
    """Display image for 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)  # pause a bit so that plots are updated


# Get a batch of training data
inputs, classes = next(iter(dataloaders['train']))

# Make a grid from batch
out = torchvision.utils.make_grid(inputs)

imshow(out, title=[class_names[x] for x in classes])

训练模型#

现在，让我们编写一个通用函数来训练模型。在这里，我们将演示：

调整学习率
保存最佳模型

在下文中，参数 scheduler 是来自 torch.optim.lr_scheduler 的学习率调度器对象。

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
    since = time.time()

    # Create a temporary directory to save training checkpoints
    with TemporaryDirectory() as tempdir:
        best_model_params_path = os.path.join(tempdir, 'best_model_params.pt')

        torch.save(model.state_dict(), best_model_params_path)
        best_acc = 0.0

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

            # Each epoch has a training and validation phase
            for phase in ['train', 'val']:
                if phase == 'train':
                    model.train()  # Set model to training mode
                else:
                    model.eval()   # Set model to evaluate mode

                running_loss = 0.0
                running_corrects = 0

                # Iterate over data.
                for inputs, labels in dataloaders[phase]:
                    inputs = inputs.to(device)
                    labels = labels.to(device)

                    # zero the parameter gradients
                    optimizer.zero_grad()

                    # forward
                    # track history if only in train
                    with torch.set_grad_enabled(phase == 'train'):
                        outputs = model(inputs)
                        _, preds = torch.max(outputs, 1)
                        loss = criterion(outputs, labels)

                        # backward + optimize only if in training phase
                        if phase == 'train':
                            loss.backward()
                            optimizer.step()

                    # statistics
                    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}')

                # deep copy the model
                if phase == 'val' and epoch_acc > best_acc:
                    best_acc = epoch_acc
                    torch.save(model.state_dict(), best_model_params_path)

            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}')

        # load best model weights
        model.load_state_dict(torch.load(best_model_params_path, weights_only=True))
    return model

可视化模型预测结果#

用于显示部分图像预测结果的通用函数

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)

微调卷积网络#

加载预训练模型并重置最后一层全连接层。

model_ft = models.resnet18(weights='IMAGENET1K_V1')
num_ftrs = model_ft.fc.in_features
# Here the size of each output sample is set to 2.
# Alternatively, it can be generalized to ``nn.Linear(num_ftrs, len(class_names))``.
model_ft.fc = nn.Linear(num_ftrs, 2)

model_ft = model_ft.to(device)

criterion = nn.CrossEntropyLoss()

# Observe that all parameters are being optimized
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

Downloading: "https://download.pytorch.org/models/resnet18-f37072fd.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/resnet18-f37072fd.pth

  0%|          | 0.00/44.7M [00:00<?, ?B/s]
 91%|█████████ | 40.5M/44.7M [00:00<00:00, 425MB/s]
100%|██████████| 44.7M/44.7M [00:00<00:00, 426MB/s]

训练与评估#

在 CPU 上大约需要 15-25 分钟。但在 GPU 上，不到一分钟即可完成。

model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                       num_epochs=25)

Epoch 0/24
----------
train Loss: 0.6229 Acc: 0.7213
val Loss: 0.2318 Acc: 0.9216

Epoch 1/24
----------
train Loss: 0.6955 Acc: 0.6885
val Loss: 0.2315 Acc: 0.9085

Epoch 2/24
----------
train Loss: 0.6665 Acc: 0.7418
val Loss: 0.2108 Acc: 0.9346

Epoch 3/24
----------
train Loss: 0.5244 Acc: 0.8115
val Loss: 0.5298 Acc: 0.8366

Epoch 4/24
----------
train Loss: 0.5970 Acc: 0.7582
val Loss: 0.5967 Acc: 0.7451

Epoch 5/24
----------
train Loss: 0.4371 Acc: 0.8361
val Loss: 0.2922 Acc: 0.8693

Epoch 6/24
----------
train Loss: 0.5674 Acc: 0.7664
val Loss: 0.3279 Acc: 0.8627

Epoch 7/24
----------
train Loss: 0.3649 Acc: 0.8566
val Loss: 0.2387 Acc: 0.9150

Epoch 8/24
----------
train Loss: 0.3884 Acc: 0.8484
val Loss: 0.2225 Acc: 0.9150

Epoch 9/24
----------
train Loss: 0.4063 Acc: 0.7992
val Loss: 0.3170 Acc: 0.8497

Epoch 10/24
----------
train Loss: 0.2582 Acc: 0.8730
val Loss: 0.2128 Acc: 0.9216

Epoch 11/24
----------
train Loss: 0.3158 Acc: 0.8730
val Loss: 0.2212 Acc: 0.9085

Epoch 12/24
----------
train Loss: 0.3275 Acc: 0.8443
val Loss: 0.2107 Acc: 0.9150

Epoch 13/24
----------
train Loss: 0.2847 Acc: 0.8730
val Loss: 0.2207 Acc: 0.9020

Epoch 14/24
----------
train Loss: 0.1782 Acc: 0.9426
val Loss: 0.1996 Acc: 0.9216

Epoch 15/24
----------
train Loss: 0.2223 Acc: 0.8934
val Loss: 0.1997 Acc: 0.9150

Epoch 16/24
----------
train Loss: 0.2765 Acc: 0.8730
val Loss: 0.2010 Acc: 0.9085

Epoch 17/24
----------
train Loss: 0.2809 Acc: 0.8607
val Loss: 0.2146 Acc: 0.9150

Epoch 18/24
----------
train Loss: 0.3457 Acc: 0.8607
val Loss: 0.2066 Acc: 0.9150

Epoch 19/24
----------
train Loss: 0.3623 Acc: 0.8402
val Loss: 0.2025 Acc: 0.9150

Epoch 20/24
----------
train Loss: 0.3433 Acc: 0.8607
val Loss: 0.2130 Acc: 0.9020

Epoch 21/24
----------
train Loss: 0.2213 Acc: 0.9221
val Loss: 0.2254 Acc: 0.9216

Epoch 22/24
----------
train Loss: 0.2951 Acc: 0.8811
val Loss: 0.2012 Acc: 0.9216

Epoch 23/24
----------
train Loss: 0.2640 Acc: 0.8893
val Loss: 0.2195 Acc: 0.9085

Epoch 24/24
----------
train Loss: 0.2251 Acc: 0.9057
val Loss: 0.2242 Acc: 0.9150

Training complete in 0m 36s
Best val Acc: 0.934641

visualize_model(model_ft)

predicted: bees, predicted: bees, predicted: ants, predicted: ants, predicted: bees, predicted: ants

卷积网络作为固定特征提取器#

在这里，我们需要冻结除了最后一层之外的所有网络。我们需要设置 requires_grad = False 来冻结参数，这样在 backward() 中就不会计算梯度。

你可以在文档此处阅读关于此内容的更多信息。

model_conv = torchvision.models.resnet18(weights='IMAGENET1K_V1')
for param in model_conv.parameters():
    param.requires_grad = False

# Parameters of newly constructed modules have requires_grad=True by default
num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 2)

model_conv = model_conv.to(device)

criterion = nn.CrossEntropyLoss()

# Observe that only parameters of final layer are being optimized as
# opposed to before.
optimizer_conv = optim.SGD(model_conv.fc.parameters(), lr=0.001, momentum=0.9)

# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)

训练与评估#

在 CPU 上，这将比上一个场景节省约一半的时间。这是预料之中的，因为网络大部分层不需要计算梯度。不过，前向传播仍然需要计算。

model_conv = train_model(model_conv, criterion, optimizer_conv,
                         exp_lr_scheduler, num_epochs=25)

Epoch 0/24
----------
train Loss: 0.5256 Acc: 0.7254
val Loss: 0.2735 Acc: 0.8954

Epoch 1/24
----------
train Loss: 0.5033 Acc: 0.7828
val Loss: 1.3175 Acc: 0.5621

Epoch 2/24
----------
train Loss: 0.7428 Acc: 0.7705
val Loss: 0.2103 Acc: 0.9281

Epoch 3/24
----------
train Loss: 0.5383 Acc: 0.7705
val Loss: 0.5245 Acc: 0.8105

Epoch 4/24
----------
train Loss: 0.5406 Acc: 0.7787
val Loss: 0.2201 Acc: 0.9412

Epoch 5/24
----------
train Loss: 0.4546 Acc: 0.8279
val Loss: 0.1917 Acc: 0.9542

Epoch 6/24
----------
train Loss: 0.4908 Acc: 0.8074
val Loss: 0.1968 Acc: 0.9346

Epoch 7/24
----------
train Loss: 0.3160 Acc: 0.8811
val Loss: 0.2180 Acc: 0.9281

Epoch 8/24
----------
train Loss: 0.3436 Acc: 0.8730
val Loss: 0.1945 Acc: 0.9346

Epoch 9/24
----------
train Loss: 0.3876 Acc: 0.8320
val Loss: 0.1944 Acc: 0.9412

Epoch 10/24
----------
train Loss: 0.2757 Acc: 0.8934
val Loss: 0.2416 Acc: 0.9085

Epoch 11/24
----------
train Loss: 0.2679 Acc: 0.8770
val Loss: 0.1863 Acc: 0.9412

Epoch 12/24
----------
train Loss: 0.3448 Acc: 0.8648
val Loss: 0.1871 Acc: 0.9412

Epoch 13/24
----------
train Loss: 0.1942 Acc: 0.9262
val Loss: 0.2183 Acc: 0.9281

Epoch 14/24
----------
train Loss: 0.3933 Acc: 0.8443
val Loss: 0.1943 Acc: 0.9412

Epoch 15/24
----------
train Loss: 0.3667 Acc: 0.8525
val Loss: 0.1995 Acc: 0.9346

Epoch 16/24
----------
train Loss: 0.3705 Acc: 0.8402
val Loss: 0.1785 Acc: 0.9412

Epoch 17/24
----------
train Loss: 0.2630 Acc: 0.8852
val Loss: 0.2075 Acc: 0.9412

Epoch 18/24
----------
train Loss: 0.3376 Acc: 0.8443
val Loss: 0.2918 Acc: 0.8824

Epoch 19/24
----------
train Loss: 0.2968 Acc: 0.8689
val Loss: 0.1928 Acc: 0.9346

Epoch 20/24
----------
train Loss: 0.3954 Acc: 0.7992
val Loss: 0.1755 Acc: 0.9412

Epoch 21/24
----------
train Loss: 0.3202 Acc: 0.8566
val Loss: 0.2193 Acc: 0.9281

Epoch 22/24
----------
train Loss: 0.3457 Acc: 0.8689
val Loss: 0.1867 Acc: 0.9281

Epoch 23/24
----------
train Loss: 0.3767 Acc: 0.8320
val Loss: 0.1817 Acc: 0.9412

Epoch 24/24
----------
train Loss: 0.4384 Acc: 0.8320
val Loss: 0.1964 Acc: 0.9281

Training complete in 0m 28s
Best val Acc: 0.954248

visualize_model(model_conv)

plt.ioff()
plt.show()

predicted: bees, predicted: ants, predicted: ants, predicted: bees, predicted: bees, predicted: ants

对自定义图像进行推理#

使用训练好的模型对自定义图像进行预测，并可视化预测的类别标签及对应的图像。

def visualize_model_predictions(model,img_path):
    was_training = model.training
    model.eval()

    img = Image.open(img_path)
    img = data_transforms['val'](img)
    img = img.unsqueeze(0)
    img = img.to(device)

    with torch.no_grad():
        outputs = model(img)
        _, preds = torch.max(outputs, 1)

        ax = plt.subplot(2,2,1)
        ax.axis('off')
        ax.set_title(f'Predicted: {class_names[preds[0]]}')
        imshow(img.cpu().data[0])

        model.train(mode=was_training)

visualize_model_predictions(
    model_conv,
    img_path='data/hymenoptera_data/val/bees/72100438_73de9f17af.jpg'
)

plt.ioff()
plt.show()

深入学习#

如果你想了解更多关于迁移学习的应用，请查看我们的计算机视觉量化迁移学习教程。

脚本总运行时间： (1 分 6.620 秒)

计算机视觉迁移学习教程#

加载数据#

可视化部分图像#

训练模型#

可视化模型预测结果#

微调卷积网络#

训练与评估#

卷积网络作为固定特征提取器#

训练与评估#

对自定义图像进行推理#

深入学习#

文档

教程

资源