PyTorch之迁移学习（Transfer Learning）

TRANSFER LEARNING TUTORIAL

• Finetuning the convnet: Instead of random initializaion, we initialize the network with a pretrained network, like the one that is trained on imagenet 1000 dataset. Rest of the training looks as usual.
• ConvNet as fixed feature extractor: Here, we will freeze the weights for all of the network except that of the final fully connected layer. This last fully connected layer is replaced with a new one with random weights and only this layer is trained.

import os
import time
import copy

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler

import numpy as np
import matplotlib.pyplot as plt

import torchvision
from torchvision import datasets, models, transforms

LOAD DATA

###
# ? normalize: mean and std values
# ? image show: [0,1] to [0,255]
###

# Data augmentation and normalization for training
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(), # range [0, 255] -> [0.0,1.0]
        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 = '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

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(0.001)
    
inputs, classes = next(iter(dataloaders['train']))

out = torchvision.utils.make_grid(inputs)

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

TRAINING THE MODEL

• Scheduling the learning rate
• Saving the best model

def train_model(model, criterion, optimizer, scheduler, 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':
                scheduler.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_sizes[phase]
            epoch_acc = running_corrects.double()/dataset_sizes[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
    

Visualizing the model predictions

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_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)

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)

exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

Train and evaluate

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

Epoch 0/24
----------
train Loss: 0.7851 Acc: 0.6270
val Loss: 1.0494 Acc: 0.6863

Epoch 1/24
----------
train Loss: 0.4664 Acc: 0.8033
val Loss: 0.2943 Acc: 0.8954

Epoch 2/24
----------
train Loss: 0.6914 Acc: 0.7254
val Loss: 0.2396 Acc: 0.9281

Epoch 3/24
----------
train Loss: 0.5356 Acc: 0.7951
val Loss: 0.3148 Acc: 0.9281

Epoch 4/24
----------
train Loss: 0.5372 Acc: 0.7910
val Loss: 0.5439 Acc: 0.8366

Epoch 5/24
----------
train Loss: 0.6590 Acc: 0.7787
val Loss: 0.3142 Acc: 0.8889

Epoch 6/24
----------
train Loss: 0.5722 Acc: 0.8033
val Loss: 0.5639 Acc: 0.8170

Epoch 7/24
----------
train Loss: 0.3940 Acc: 0.8402
val Loss: 0.2142 Acc: 0.9085

Epoch 8/24
----------
train Loss: 0.3693 Acc: 0.8484
val Loss: 0.2529 Acc: 0.8758

Epoch 9/24
----------
train Loss: 0.2801 Acc: 0.8852
val Loss: 0.2550 Acc: 0.8758

Epoch 10/24
----------
train Loss: 0.2674 Acc: 0.8934
val Loss: 0.2333 Acc: 0.9085

Epoch 11/24
----------
train Loss: 0.2915 Acc: 0.8689
val Loss: 0.2292 Acc: 0.9085

Epoch 12/24
----------
train Loss: 0.3659 Acc: 0.8443
val Loss: 0.2254 Acc: 0.9085

Epoch 13/24
----------
train Loss: 0.3359 Acc: 0.8525
val Loss: 0.1940 Acc: 0.9150

Epoch 14/24
----------
train Loss: 0.2926 Acc: 0.8689
val Loss: 0.2075 Acc: 0.8954

Epoch 15/24
----------
train Loss: 0.2606 Acc: 0.8934
val Loss: 0.2457 Acc: 0.9085

Epoch 16/24
----------
train Loss: 0.2715 Acc: 0.8770
val Loss: 0.2307 Acc: 0.9020

Epoch 17/24
----------
train Loss: 0.2026 Acc: 0.8975
val Loss: 0.2478 Acc: 0.9020

Epoch 18/24
----------
train Loss: 0.2696 Acc: 0.8934
val Loss: 0.2385 Acc: 0.8954

Epoch 19/24
----------
train Loss: 0.2277 Acc: 0.9016
val Loss: 0.1987 Acc: 0.9085

Epoch 20/24
----------
train Loss: 0.2901 Acc: 0.9016
val Loss: 0.2024 Acc: 0.9020

Epoch 21/24
----------
train Loss: 0.2981 Acc: 0.8811
val Loss: 0.2100 Acc: 0.8889

Epoch 22/24
----------
train Loss: 0.2462 Acc: 0.8893
val Loss: 0.2201 Acc: 0.9020

Epoch 23/24
----------
train Loss: 0.3056 Acc: 0.8730
val Loss: 0.1911 Acc: 0.9150

Epoch 24/24
----------
train Loss: 0.2561 Acc: 0.9057
val Loss: 0.2074 Acc: 0.9216

Training complete in 3m 57s
Best val Acc: 0.928105

visualize_model(model_ft)

CONVNET AS FIXED FEATURE EXTRACTOR

model_conv = torchvision.models.resnet18(pretrained=True)

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
# opoosed 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)

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

Epoch 0/24
----------
train Loss: 0.6007 Acc: 0.6967
val Loss: 0.3234 Acc: 0.8627

Epoch 1/24
----------
train Loss: 0.4813 Acc: 0.7664
val Loss: 0.2608 Acc: 0.8954

Epoch 2/24
----------
train Loss: 0.4790 Acc: 0.7992
val Loss: 0.2120 Acc: 0.9477

Epoch 3/24
----------
train Loss: 0.5683 Acc: 0.7869
val Loss: 0.3591 Acc: 0.8497

Epoch 4/24
----------
train Loss: 0.3770 Acc: 0.8402
val Loss: 0.2527 Acc: 0.9020

Epoch 5/24
----------
train Loss: 0.4000 Acc: 0.8279
val Loss: 0.1915 Acc: 0.9346

Epoch 6/24
----------
train Loss: 0.3457 Acc: 0.8402
val Loss: 0.2083 Acc: 0.9477

Epoch 7/24
----------
train Loss: 0.3756 Acc: 0.8443
val Loss: 0.2210 Acc: 0.9346

Epoch 8/24
----------
train Loss: 0.3511 Acc: 0.8607
val Loss: 0.2047 Acc: 0.9412

Epoch 9/24
----------
train Loss: 0.4228 Acc: 0.8361
val Loss: 0.2112 Acc: 0.9412

Epoch 10/24
----------
train Loss: 0.3002 Acc: 0.8648
val Loss: 0.1862 Acc: 0.9477

Epoch 11/24
----------
train Loss: 0.4303 Acc: 0.8320
val Loss: 0.2357 Acc: 0.9216

Epoch 12/24
----------
train Loss: 0.3457 Acc: 0.8320
val Loss: 0.1871 Acc: 0.9477

Epoch 13/24
----------
train Loss: 0.4396 Acc: 0.8033
val Loss: 0.2109 Acc: 0.9346

Epoch 14/24
----------
train Loss: 0.2644 Acc: 0.8852
val Loss: 0.2328 Acc: 0.9412

Epoch 15/24
----------
train Loss: 0.3323 Acc: 0.8689
val Loss: 0.2145 Acc: 0.9346

Epoch 16/24
----------
train Loss: 0.3899 Acc: 0.7951
val Loss: 0.2646 Acc: 0.9150

Epoch 17/24
----------
train Loss: 0.4303 Acc: 0.8074
val Loss: 0.2338 Acc: 0.9281

Epoch 18/24
----------
train Loss: 0.3758 Acc: 0.8607
val Loss: 0.2158 Acc: 0.9346

Epoch 19/24
----------
train Loss: 0.2428 Acc: 0.9303
val Loss: 0.2132 Acc: 0.9281

Epoch 20/24
----------
train Loss: 0.3515 Acc: 0.8566
val Loss: 0.2062 Acc: 0.9412

Epoch 21/24
----------
train Loss: 0.3181 Acc: 0.8689
val Loss: 0.2119 Acc: 0.9281

Epoch 22/24
----------
train Loss: 0.3594 Acc: 0.8156
val Loss: 0.2083 Acc: 0.9412

Epoch 23/24
----------
train Loss: 0.3170 Acc: 0.8648
val Loss: 0.1905 Acc: 0.9346

Epoch 24/24
----------
train Loss: 0.3265 Acc: 0.8607
val Loss: 0.2156 Acc: 0.9281

Training complete in 2m 52s
Best val Acc: 0.947712

visualize_model(model_conv)