来自这里。
在本教程中,你将学习如何使用迁移学习来训练你的网络。在cs231n notes你可以了解更多关于迁移学习的知识。
在实践中,很少有人从头开始训练整个卷积网络(使用随机初始化),因为拥有足够大小的数据集相对较少。相反,通常在非常大的数据集(例如ImageNet,它包含120万幅、1000个类别的图像)上对ConvNet进行预训练,然后使用ConvNet作为初始化或固定的特征提取器来执行感兴趣的任务。
两个主要的迁移学习的场景如下:
- Finetuning the convert:与随机初始化不同,我们使用一个预训练的网络初始化网络,就像在imagenet 1000 dataset上训练的网络一样。其余的训练看起来和往常一样。
- ConvNet as fixed feature extractor:在这里,我们将冻结所有网络的权重,除了最后的全连接层。最后一个全连接层被替换为一个具有随机权重的新层,并且只训练这一层。
#!/usr/bin/env python3 # License: BSD # Author: Sasank Chilamkurthy from __future__ import print_function,division import torch import torch.nn as nn import torch.optim as optim import numpy as np import torchvision from torchvision import datasets,models,transforms import matplotlib.pyplot as plt import time import os import copy plt.ion() # 交互模式
我们使用torchvision
和torch.utils.data
包来导入数据。
我们今天要解决的问题是训练一个模型来区分蚂蚁和蜜蜂。我们有蚂蚁和蜜蜂的训练图像各120张。每一类有75张验证图片。通常,如果是从零开始训练,这是一个非常小的数据集。因为我们要使用迁移学习,所以我们的例子应该具有很好地代表性。
这个数据集是一个非常小的图像子集。
你可以从这里下载数据并解压到当前目录。
# 训练数据的扩充及标准化 # 只进行标准化验证 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_size = {x:len(image_datasets[x]) for x in ['train','val']} class_name = image_datasets['train'].classes device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
为了理解数据扩充,我们可视化一些训练图像。
def imshow(inp, title=None): 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(10) # 暂停一会,以便更新绘图 # 获取一批训练数据 inputs, classes = next(iter(dataloaders['train'])) # 从批处理中生成网格 out = torchvision.utils.make_grid(inputs) imshow(out, title=[class_name[x] for x in classes])
现在我们来实现一个通用函数来训练一个模型。在这个函数中,我们将:
- 调整学习率
- 保存最优模型
下面例子中,参数schedule
是来自torch.optim.lr_scheduler
的LR调度对象。
def train_model(model, criterion, optimizer, schduler, num_epochs=25): since = time.time() best_model_wts = copy.deepcopy(model.state_dict()) best_acc = 0.0 for epoch in range(num_epochs): print('Epoch {}/{}'.format(epoch, num_epochs-1)) print('-'*10) for phase in ['train', 'val']: if phase == 'train': schduler.step() 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() # 统计 running_loss += loss.item() * inputs.size(0) running_corrects += torch.sum(preds == labels.data) epoch_loss = running_loss / dataset_size[phase] epoch_acc = running_corrects.double() / dataset_size[phase] print('{} Loss: {:.4f} Acc: {:.4f}'.format( phase, epoch_loss, epoch_acc)) # 很拷贝模型 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('Training complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) print('Best val Acc: {:4f}'.format(best_acc)) # 导入最优模型权重 model.load_state_dict(best_model_wts) 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('predicted: {}'.format(class_name[preds[j]])) imshow(inputs.cpu().data[j]) if images_so_far == num_images: model.train(mode=was_training) return model.train(mode=was_training)
Finetuning the convnet
加载预处理的模型和重置最后的全连接层:
model_ft = models.resnet18(pretrained=True) num_ftrs = model_ft.fc.in_features model_ft.fc = nn.Linear(num_ftrs, 2) model_ft = model_ft.to(device) criterion = nn.CrossEntropyLoss() # 优化所有参数 optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9) # 没7次,学习率衰减0.1 exp_lr_scheduler = torch.optim.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, num_epochs=25)
visualize_model(model_ft)
现在,我们冻结除最后一层外的所有网络。我们需要设置requires_grad=False
来冻结参数,这样调用backward()
时不计算梯度。
你可以从这篇文档中了解更多。
model_conv = models.resnet18(pretrained=True) for param in model_conv.parameters(): param.requires_grad = False # 新构造模块的参数默认requires_grad=True num_ftrs = model_conv.fc.in_features model_conv.fc = nn.Linear(num_ftrs, 2) model_conv = model_conv.to(device) criterion = nn.CrossEntropyLoss() # 优化所有参数 optimizer_ft = optim.SGD(model_conv.parameters(), lr=0.001, momentum=0.9) # 没7次,学习率衰减0.1 exp_lr_scheduler = torch.optim.lr_scheduler.StepLR( optimizer_ft, step_size=7, gamma=0.1) model_conv = train_model(model_conv, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=25)
visualize_model(model_conv) plt.ioff() plt.show()