迁移学习教程

作者：Sasank Chilamkurthy

译者：DrDavidS

校验：DrDavidS

在本教程中，您将学习如何使用迁移学习训练网络。你可以在 cs231n笔记中阅读更多关于迁移学习的内容。

引用笔记，

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

如下是两个主要的迁移学习场景：

  • 微调卷积神经网络 我们使用预训练网络来初始化网络，而不是随机初始化，比如一个已经在imagenet 1000数据集上训练好的网络一样。其余训练和往常一样。
  • 将卷积神经网络作为固定特征提取器 ：在这里，我们将冻结除最终全连接层之外的整个网络的权重。最后一个全连接层被替换为具有随机权重的新层，并且仅训练该层。
  1. # License: BSD
  2. # Author: Sasank Chilamkurthy
  4. from __future__ import print_function, division
  6. import torch
  7. import torch.nn as nn
  8. import torch.optim as optim
  9. from torch.optim import lr_scheduler
  10. import numpy as np
  11. import torchvision
  12. from torchvision import datasets, models, transforms
  13. import matplotlib.pyplot as plt
  14. import time
  15. import os
  16. import copy
  18. plt.ion()   # interactive mode

加载数据

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

今天，我们要解决的问题是训练一个模型来对蚂蚁和蜜蜂进行分类。我们蚂蚁和蜜蜂分别准备了大约120个训练图像，并且每类还有75个验证图像。通常，如果从头开始训练，这是一个非常小的数据集。由于我们正在使用迁移学习，我们应该能够合理地进行泛化。

该数据集是imagenet的一个很小的子集。

注意

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

  1. # Data augmentation and normalization for training
  2. # Just normalization for validation
  3. data_transforms = {
  4.     'train': transforms.Compose([
  5.         transforms.RandomResizedCrop(224),
  6.         transforms.RandomHorizontalFlip(),
  7.         transforms.ToTensor(),
  8.         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  9.     ]),
  10.     'val': transforms.Compose([
  11.         transforms.Resize(256),
  12.         transforms.CenterCrop(224),
  13.         transforms.ToTensor(),
  14.         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  15.     ]),
  16. }
  18. data_dir = 'data/hymenoptera_data'
  19. image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
  20.                                           data_transforms[x])
  21.                   for x in ['train', 'val']}
  22. dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=4,
  23.                                              shuffle=True, num_workers=4)
  24.               for x in ['train', 'val']}
  25. dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
  26. class_names = image_datasets['train'].classes
  28. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

可视化一些图像

让我们通过可视化一些训练图像，来理解什么是数据增强。

  1. def imshow(inp, title=None):
  2.     """Imshow for Tensor."""
  3.     inp = inp.numpy().transpose((1, 2, 0))
  4.     mean = np.array([0.485, 0.456, 0.406])
  5.     std = np.array([0.229, 0.224, 0.225])
  6.     inp = std * inp + mean
  7.     inp = np.clip(inp, 0, 1)
  8.     plt.imshow(inp)
  9.     if title is not None:
  10.         plt.title(title)
  11.     plt.pause(0.001)  # pause a bit so that plots are updated
  14. # Get a batch of training data
  15. inputs, classes = next(iter(dataloaders['train']))
  17. # Make a grid from batch
  18. out = torchvision.utils.make_grid(inputs)
  20. imshow(out, title=[class_names[x] for x in classes])

img/sphx_glr_transfer_learning_tutorial_001.png

训练模型

现在, 让我们编写一个通用函数来训练一个模型。这里, 我们将会举例说明:

  • 调整学习率
  • 保存最好的模型

下面函数中, scheduler 参数是 torch.optim.lr_scheduler 中的学习率调整(LR scheduler）对象.

  1. def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
  2.     since = time.time()
  4.     best_model_wts = copy.deepcopy(model.state_dict())
  5.     best_acc = 0.0
  7.     for epoch in range(num_epochs):
  8.         print('Epoch {}/{}'.format(epoch, num_epochs - 1))
  9.         print('-' * 10)
  11.         # Each epoch has a training and validation phase
  12.         for phase in ['train', 'val']:
  13.             if phase == 'train':
  14.                 model.train()  # Set model to training mode
  15.             else:
  16.                 model.eval()   # Set model to evaluate mode
  18.             running_loss = 0.0
  19.             running_corrects = 0
  21.             # Iterate over data.
  22.             for inputs, labels in dataloaders[phase]:
  23.                 inputs = inputs.to(device)
  24.                 labels = labels.to(device)
  26.                 # zero the parameter gradients
  27.                 optimizer.zero_grad()
  29.                 # forward
  30.                 # track history if only in train
  31.                 with torch.set_grad_enabled(phase == 'train'):
  32.                     outputs = model(inputs)
  33.                     _, preds = torch.max(outputs, 1)
  34.                     loss = criterion(outputs, labels)
  36.                     # backward + optimize only if in training phase
  37.                     if phase == 'train':
  38.                         loss.backward()
  39.                         optimizer.step()
  41.                 # statistics
  42.                 running_loss += loss.item() * inputs.size(0)
  43.                 running_corrects += torch.sum(preds == labels.data)
  44.             if phase == 'train':
  45.                 scheduler.step()
  47.             epoch_loss = running_loss / dataset_sizes[phase]
  48.             epoch_acc = running_corrects.double() / dataset_sizes[phase]
  50.             print('{} Loss: {:.4f} Acc: {:.4f}'.format(
  51.                 phase, epoch_loss, epoch_acc))
  53.             # deep copy the model
  54.             if phase == 'val' and epoch_acc > best_acc:
  55.                 best_acc = epoch_acc
  56.                 best_model_wts = copy.deepcopy(model.state_dict())
  58.         print()
  60.     time_elapsed = time.time() - since
  61.     print('Training complete in {:.0f}m {:.0f}s'.format(
  62.         time_elapsed // 60, time_elapsed % 60))
  63.     print('Best val Acc: {:4f}'.format(best_acc))
  65.     # load best model weights
  66.     model.load_state_dict(best_model_wts)
  67.     return model

模型预测的可视化

用于显示少量预测图像的通用函数

  1. def visualize_model(model, num_images=6):
  2.     was_training = model.training
  3.     model.eval()
  4.     images_so_far = 0
  5.     fig = plt.figure()
  7.     with torch.no_grad():
  8.         for i, (inputs, labels) in enumerate(dataloaders['val']):
  9.             inputs = inputs.to(device)
  10.             labels = labels.to(device)
  12.             outputs = model(inputs)
  13.             _, preds = torch.max(outputs, 1)
  15.             for j in range(inputs.size()[0]):
  16.                 images_so_far += 1
  17.                 ax = plt.subplot(num_images//2, 2, images_so_far)
  18.                 ax.axis('off')
  19.                 ax.set_title('predicted: {}'.format(class_names[preds[j]]))
  20.                 imshow(inputs.cpu().data[j])
  22.                 if images_so_far == num_images:
  23.                     model.train(mode=was_training)
  24.                     return
  25.         model.train(mode=was_training)

微调卷积神经网络

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

  1. model_ft = models.resnet18(pretrained=True)
  2. num_ftrs = model_ft.fc.in_features
  3. # Here the size of each output sample is set to 2.
  4. # Alternatively, it can be generalized to nn.Linear(num_ftrs, len(class_names)).
  5. model_ft.fc = nn.Linear(num_ftrs, 2)
  7. model_ft = model_ft.to(device)
  9. criterion = nn.CrossEntropyLoss()
  11. # Observe that all parameters are being optimized
  12. optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
  14. # Decay LR by a factor of 0.1 every 7 epochs
  15. exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

训练与评价

在CPU上训练需要大约15-25分钟。但是在GPU上，它只需不到一分钟。

  1. model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
  2.                        num_epochs=25)

输出：

  1. Epoch 0/24
  2. ----------
  3. train Loss: 0.6751 Acc: 0.7049
  4. val Loss: 0.1834 Acc: 0.9346
  6. Epoch 1/24
  7. ----------
  8. train Loss: 0.5892 Acc: 0.7746
  9. val Loss: 1.0048 Acc: 0.6667
  11. Epoch 2/24
  12. ----------
  13. train Loss: 0.6568 Acc: 0.7459
  14. val Loss: 0.6047 Acc: 0.8366
  16. Epoch 3/24
  17. ----------
  18. train Loss: 0.4196 Acc: 0.8320
  19. val Loss: 0.4388 Acc: 0.8562
  21. Epoch 4/24
  22. ----------
  23. train Loss: 0.5883 Acc: 0.8033
  24. val Loss: 0.4013 Acc: 0.8889
  26. Epoch 5/24
  27. ----------
  28. train Loss: 0.6684 Acc: 0.7705
  29. val Loss: 0.2666 Acc: 0.9412
  31. Epoch 6/24
  32. ----------
  33. train Loss: 0.5308 Acc: 0.7787
  34. val Loss: 0.4803 Acc: 0.8693
  36. Epoch 7/24
  37. ----------
  38. train Loss: 0.3464 Acc: 0.8566
  39. val Loss: 0.2385 Acc: 0.8954
  41. Epoch 8/24
  42. ----------
  43. train Loss: 0.4586 Acc: 0.7910
  44. val Loss: 0.2064 Acc: 0.9020
  46. Epoch 9/24
  47. ----------
  48. train Loss: 0.3438 Acc: 0.8402
  49. val Loss: 0.2336 Acc: 0.9020
  51. Epoch 10/24
  52. ----------
  53. train Loss: 0.2405 Acc: 0.9016
  54. val Loss: 0.1866 Acc: 0.9346
  56. Epoch 11/24
  57. ----------
  58. train Loss: 0.2335 Acc: 0.8852
  59. val Loss: 0.2152 Acc: 0.9216
  61. Epoch 12/24
  62. ----------
  63. train Loss: 0.3441 Acc: 0.8402
  64. val Loss: 0.2298 Acc: 0.9020
  66. Epoch 13/24
  67. ----------
  68. train Loss: 0.2513 Acc: 0.9098
  69. val Loss: 0.2204 Acc: 0.9020
  71. Epoch 14/24
  72. ----------
  73. train Loss: 0.2745 Acc: 0.8934
  74. val Loss: 0.2439 Acc: 0.8889
  76. Epoch 15/24
  77. ----------
  78. train Loss: 0.2978 Acc: 0.8607
  79. val Loss: 0.2817 Acc: 0.8497
  81. Epoch 16/24
  82. ----------
  83. train Loss: 0.2560 Acc: 0.8975
  84. val Loss: 0.1933 Acc: 0.9281
  86. Epoch 17/24
  87. ----------
  88. train Loss: 0.2326 Acc: 0.9098
  89. val Loss: 0.2176 Acc: 0.9085
  91. Epoch 18/24
  92. ----------
  93. train Loss: 0.2274 Acc: 0.9016
  94. val Loss: 0.2084 Acc: 0.9346
  96. Epoch 19/24
  97. ----------
  98. train Loss: 0.3091 Acc: 0.8689
  99. val Loss: 0.2270 Acc: 0.9150
  101. Epoch 20/24
  102. ----------
  103. train Loss: 0.2540 Acc: 0.8975
  104. val Loss: 0.1957 Acc: 0.9216
  106. Epoch 21/24
  107. ----------
  108. train Loss: 0.3203 Acc: 0.8648
  109. val Loss: 0.1969 Acc: 0.9216
  111. Epoch 22/24
  112. ----------
  113. train Loss: 0.3048 Acc: 0.8443
  114. val Loss: 0.1981 Acc: 0.9346
  116. Epoch 23/24
  117. ----------
  118. train Loss: 0.2526 Acc: 0.9016
  119. val Loss: 0.2415 Acc: 0.8889
  121. Epoch 24/24
  122. ----------
  123. train Loss: 0.3041 Acc: 0.8689
  124. val Loss: 0.1894 Acc: 0.9346
  126. Training complete in 1m 7s
  127. Best val Acc: 0.941176
  131. visualize_model(model_ft)

img/sphx_glr_transfer_learning_tutorial_002.png

将卷积神经网络为固定特征提取器

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

您可以在此处的文档中阅读更多相关信息。

  1. model_conv = torchvision.models.resnet18(pretrained=True)
  2. for param in model_conv.parameters():
  3.     param.requires_grad = False
  5. # Parameters of newly constructed modules have requires_grad=True by default
  6. num_ftrs = model_conv.fc.in_features
  7. model_conv.fc = nn.Linear(num_ftrs, 2)
  9. model_conv = model_conv.to(device)
  11. criterion = nn.CrossEntropyLoss()
  13. # Observe that only parameters of final layer are being optimized as
  14. # opposed to before.
  15. optimizer_conv = optim.SGD(model_conv.fc.parameters(), lr=0.001, momentum=0.9)
  17. # Decay LR by a factor of 0.1 every 7 epochs
  18. exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)

训练与评价

在CPU上，与前一个场景相比，大概只花费一半的时间。这在预料之中，因为不需要为绝大多数网络计算梯度。当然，我们还是需要计算前向传播。

  1. model_conv = train_model(model_conv, criterion, optimizer_conv,
  2.                          exp_lr_scheduler, num_epochs=25)

Out:

  1. Epoch 0/24
  2. ----------
  3. train Loss: 0.6073 Acc: 0.6598
  4. val Loss: 0.2511 Acc: 0.8954
  6. Epoch 1/24
  7. ----------
  8. train Loss: 0.5457 Acc: 0.7459
  9. val Loss: 0.5169 Acc: 0.7647
  11. Epoch 2/24
  12. ----------
  13. train Loss: 0.4023 Acc: 0.8320
  14. val Loss: 0.2361 Acc: 0.9150
  16. Epoch 3/24
  17. ----------
  18. train Loss: 0.5150 Acc: 0.7869
  19. val Loss: 0.5423 Acc: 0.8039
  21. Epoch 4/24
  22. ----------
  23. train Loss: 0.4142 Acc: 0.8115
  24. val Loss: 0.2257 Acc: 0.9216
  26. Epoch 5/24
  27. ----------
  28. train Loss: 0.6364 Acc: 0.7418
  29. val Loss: 0.3133 Acc: 0.8889
  31. Epoch 6/24
  32. ----------
  33. train Loss: 0.5543 Acc: 0.7664
  34. val Loss: 0.1959 Acc: 0.9412
  36. Epoch 7/24
  37. ----------
  38. train Loss: 0.3552 Acc: 0.8443
  39. val Loss: 0.2013 Acc: 0.9477
  41. Epoch 8/24
  42. ----------
  43. train Loss: 0.3538 Acc: 0.8525
  44. val Loss: 0.1825 Acc: 0.9542
  46. Epoch 9/24
  47. ----------
  48. train Loss: 0.3954 Acc: 0.8402
  49. val Loss: 0.1959 Acc: 0.9477
  51. Epoch 10/24
  52. ----------
  53. train Loss: 0.3615 Acc: 0.8443
  54. val Loss: 0.1779 Acc: 0.9542
  56. Epoch 11/24
  57. ----------
  58. train Loss: 0.3951 Acc: 0.8320
  59. val Loss: 0.1730 Acc: 0.9542
  61. Epoch 12/24
  62. ----------
  63. train Loss: 0.4111 Acc: 0.8156
  64. val Loss: 0.2573 Acc: 0.9150
  66. Epoch 13/24
  67. ----------
  68. train Loss: 0.3073 Acc: 0.8525
  69. val Loss: 0.1901 Acc: 0.9477
  71. Epoch 14/24
  72. ----------
  73. train Loss: 0.3288 Acc: 0.8279
  74. val Loss: 0.2114 Acc: 0.9346
  76. Epoch 15/24
  77. ----------
  78. train Loss: 0.3472 Acc: 0.8525
  79. val Loss: 0.1989 Acc: 0.9412
  81. Epoch 16/24
  82. ----------
  83. train Loss: 0.3309 Acc: 0.8689
  84. val Loss: 0.1757 Acc: 0.9412
  86. Epoch 17/24
  87. ----------
  88. train Loss: 0.3963 Acc: 0.8197
  89. val Loss: 0.1881 Acc: 0.9608
  91. Epoch 18/24
  92. ----------
  93. train Loss: 0.3332 Acc: 0.8484
  94. val Loss: 0.2175 Acc: 0.9412
  96. Epoch 19/24
  97. ----------
  98. train Loss: 0.3419 Acc: 0.8320
  99. val Loss: 0.1932 Acc: 0.9412
  101. Epoch 20/24
  102. ----------
  103. train Loss: 0.3471 Acc: 0.8689
  104. val Loss: 0.1851 Acc: 0.9477
  106. Epoch 21/24
  107. ----------
  108. train Loss: 0.2843 Acc: 0.8811
  109. val Loss: 0.1772 Acc: 0.9477
  111. Epoch 22/24
  112. ----------
  113. train Loss: 0.4024 Acc: 0.8402
  114. val Loss: 0.1818 Acc: 0.9542
  116. Epoch 23/24
  117. ----------
  118. train Loss: 0.2409 Acc: 0.8975
  119. val Loss: 0.2211 Acc: 0.9346
  121. Epoch 24/24
  122. ----------
  123. train Loss: 0.3838 Acc: 0.8238
  124. val Loss: 0.1918 Acc: 0.9412
  126. Training complete in 0m 34s
  127. Best val Acc: 0.960784
  131. visualize_model(model_conv)
  133. plt.ioff()
  134. plt.show()

img/sphx_glr_transfer_learning_tutorial_003.png

脚本的总运行时间： (1分钟53.655秒）

由Sphinx-Gallery生成的图库