【小白学习PyTorch教程】十二、迁移学习:微调VGG19实现图像分类

2022-08-18 09:29:50 浏览数 (2)

「@Author:Runsen」

前言:迁移学习就是利用数据、任务或模型之间的相似性,将在旧的领域学习过或训练好的模型,应用于新的领域这样的一个过程。从这段定义里面,我们可以窥见迁移学习的关键点所在,即新的任务与旧的任务在数据、任务和模型之间的相似性。

假设有两个任务系统A和B,任务A拥有海量的数据资源且已训练好,但并不是我们的目标任务,任务B是我们的目标任务,但数据量少且极为珍贵,这种场景便是典型的迁移学习的应用场景

接下来在博客中,我们将学习如何将迁移学习与 PyTorch 结合使用。

在这个迁移学习 PyTorch 图像二分类Vgg19 示例中,数据来源:https://www.kaggle.com/pmigdal/alien-vs-predator-images/home

这是我在kaggle找到的关于迁移学习的入门案例

1) 加载数据

第一步是加载数据并对图像进行一些转换,使其符合网络要求。

使用 torchvision.dataset ,在文件夹中加载数据。该模块将在文件夹中迭代以拆分数据以进行训练和验证。

转换过程进行基本的图片处理操作。

将从中心裁剪图像,执行水平翻转,归一化,最后使用将其转换为张量。

代码语言:javascript复制
import os
import time
import torch
import torchvision
from torchvision import datasets, models, transforms
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

data_dir = "alien_pred"
input_shape = 224
mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]

#data transformation
data_transforms = {
   'train': transforms.Compose([
       transforms.CenterCrop(input_shape),
       transforms.ToTensor(),
       transforms.Normalize(mean, std)
   ]),
   'validation': transforms.Compose([
       transforms.CenterCrop(input_shape),
       transforms.ToTensor(),
       transforms.Normalize(mean, std)
   ]),
}

image_datasets = {
   x: datasets.ImageFolder(
       os.path.join(data_dir, x),
       transform=data_transforms[x]
   )
   for x in ['train', 'validation']
}

dataloaders = {
   x: torch.utils.data.DataLoader(
       image_datasets[x], batch_size=32,
       shuffle=True, num_workers=4
   )
   for x in ['train', 'validation']
}

dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'validation']}

print(dataset_sizes)
# {'train': 694, 'validation': 200}

class_names = image_datasets['train'].classes

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

可视化 PyTorch 迁移学习数据集。可视化过程将从训练数据加载器和标签中获取下一批图像,并用 matplot 显示它。

代码语言:javascript复制
images, labels = next(iter(dataloaders['train']))



rows = 4
columns = 4
fig=plt.figure(figsize=(15,15))


for i in range(16):
   fig.add_subplot(rows, columns, i 1)
   plt.title(class_names[labels[i]])
   img = images[i].numpy().transpose((1, 2, 0))
   img = std * img   mean
   plt.imshow(img)
plt.show()

2) 定义模型

VGG19有两个部分,分别是VGG19.features和VGG19.classifier。

  • vgg19.features有卷积层和池化层
  • vgg19.features有三个线性层,最后是softmax分类器

下面将使用 torchvision.models 加载 VGG19,并将预训练权重设置为 True之后,将冻结层,使这些层不可训练。

对 Linear 层修改最后一层,以满足我们 2 个类的需求。

也可以将 CrossEntropyLoss 用于多类损失函数,对于优化器,使用学习率为 0.0001 和动量为 0.9 的 SGD,如下面的 PyTorch 迁移学习示例所示。

代码语言:javascript复制
##加载基于VGG19的模型
vgg_based = torchvision.models.vgg19(pretrained=True) 

for param in vgg_based.parameters():
   param.requires_grad = False

#修改最后一层
number_features = vgg_based.classifier[6].in_features 
features = list(vgg_based.classifier.children())[:-1] # 移除最后一层
features.extend([torch.nn.Linear(number_features, len(class_names))]) 
vgg_based.classifier = torch.nn.Sequential(*features) 

vgg_based = vgg_based.to(device) 

print(vgg_based)

criterion = torch.nn.CrossEntropyLoss()
optimizer_ft = optim.SGD(vgg_based.parameters(), lr=0.001, momentum=0.9)

vgg_based输出如下

代码语言:javascript复制
VGG(
  (features): Sequential(
 (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (1): ReLU(inplace)
 (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (3): ReLU(inplace)
 (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
 (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (6): ReLU(inplace)
 (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (8): ReLU(inplace)
 (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
 (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (11): ReLU(inplace)
 (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (13): ReLU(inplace)
 (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (15): ReLU(inplace)
 (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (17): ReLU(inplace)
 (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
 (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (20): ReLU(inplace)
 (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (22): ReLU(inplace)
 (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (24): ReLU(inplace)
 (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (26): ReLU(inplace)
 (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
 (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (29): ReLU(inplace)
 (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (31): ReLU(inplace)
 (32): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (33): ReLU(inplace)
 (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (35): ReLU(inplace)
 (36): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (classifier): Sequential(
 (0): Linear(in_features=25088, out_features=4096, bias=True)
 (1): ReLU(inplace)
 (2): Dropout(p=0.5)
 (3): Linear(in_features=4096, out_features=4096, bias=True)
 (4): ReLU(inplace)
 (5): Dropout(p=0.5)
 (6): Linear(in_features=4096, out_features=2, bias=True)
  )
)

3) 训练模型

下面使用 PyTorch 中的一些功能来帮助我们训练和评估我们的模型。

代码语言:javascript复制
def train_model(model, criterion, optimizer, num_epochs=25):
   since = time.time()

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

        # 迭代数据
       train_loss = 0

       # Iterate over data.
       for i, data in enumerate(dataloaders['train']):
           inputs , labels = data
           inputs = inputs.to(device)
           labels = labels.to(device)

           optimizer.zero_grad()
          
           with torch.set_grad_enabled(True):
               outputs  = model(inputs)
               loss = criterion(outputs, labels)

           loss.backward()
           optimizer.step()

           train_loss  = loss.item() * inputs.size(0)

           print('{} Loss: {:.4f}'.format(
               'train', train_loss / dataset_sizes['train']))
          
   time_elapsed = time.time() - since
   print('Training complete in {:.0f}m {:.0f}s'.format(
       time_elapsed // 60, time_elapsed % 60))

   return model


最后, epoch 数设置为 25 开始我们的模型训练过程,并在训练过程结束后进行评估。

在每个训练步骤中,模型接受输入并预测输出。之后预测输出将传递给计算损失。然后损失将执行反向传播来计算得到梯度,最后计算权重并使用 autograd 不断的优化参数。

代码语言:javascript复制
vgg_based = train_model(vgg_based, criterion, optimizer_ft, num_epochs=25)

4) 测试模型

在可视化模型中,将训练好的模型,使用一批图像进行测试和预测标签

代码语言:javascript复制
def visualize_model(model, num_images=6):
   was_training = model.training
   model.eval()
   images_so_far = 0
   fig=plt.figure(figsize=(15,15))


   with torch.no_grad():
       for i, (inputs, labels) in enumerate(dataloaders['validation']):
           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('predicted: {} truth: {}'.format(class_names[preds[j]], class_names[labels[j]]))
               img = inputs.cpu().data[j].numpy().transpose((1, 2, 0))
               img = std * img   mean
               ax.imshow(img)

               if images_so_far == num_images:
                   model.train(mode=was_training)
                   return
       model.train(mode=was_training)
visualize_model(vgg_based)
plt.show()

0 人点赞