Pytorch Transfer learning

标签（空格分隔）： Pytorch

---

These two major transfer learning scenarios looks as follows:

1. 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.
2. 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

# License: BSD# Author: Sasank Chilamkurthyfrom __future__ import print_function, divisionimport torchimport torch.nn as nnimport torch.optim as optimfrom torch.optim import lr_schedulerfrom torch.autograd import Variableimport numpy as npimport torchvisionfrom torchvision import datasets, models, transformsimport matplotlib.pyplot as pltimport timeimport osplt.ion()   # interactive mode

Load Data

# Data augmentation and normalization for training# Just normalization for validationdata_transforms = {    'train': transforms.Compose([        transforms.RandomSizedCrop(224),        transforms.RandomHorizontalFlip(),        transforms.ToTensor(),        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])    ]),    'val': transforms.Compose([        transforms.Scale(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']}dataloders = {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'].classesuse_gpu = torch.cuda.is_available()

Visualize a few images

def imshow(inp, title=None):    """Imshow 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 datainputs, classes = next(iter(dataloders['train']))# Make a grid from batchout = torchvision.utils.make_grid(inputs)imshow(out, title=[class_names[x] for x in classes])

Trainning the model

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):    since = time.time()    best_model_wts = model.state_dict()    best_acc = 0.0    for epoch in range(num_epochs):        print('Epoch {}/{}'.format(epoch, num_epochs - 1))        print('-' * 10)        # Each epoch has a training and validation phase        for phase in ['train', 'val']:            if phase == 'train':                scheduler.step()                model.train(True)  # Set model to training mode            else:                model.train(False)  # Set model to evaluate mode            running_loss = 0.0            running_corrects = 0            # Iterate over data.            for data in dataloders[phase]:                # get the inputs                inputs, labels = data                # wrap them in Variable                if use_gpu:                    inputs = Variable(inputs.cuda())                    labels = Variable(labels.cuda())                else:                    inputs, labels = Variable(inputs), Variable(labels)                # zero the parameter gradients                optimizer.zero_grad()                # forward                outputs = model(inputs)                _, preds = torch.max(outputs.data, 1)                loss = criterion(outputs, labels)                # backward + optimize only if in training phase                if phase == 'train':                    loss.backward()                    optimizer.step()                # statistics                running_loss += loss.data[0]                running_corrects += torch.sum(preds == labels.data)            epoch_loss = running_loss / dataset_sizes[phase]            epoch_acc = running_corrects / dataset_sizes[phase]            print('{} Loss: {:.4f} Acc: {:.4f}'.format(                phase, epoch_loss, epoch_acc))            # deep copy the model            if phase == 'val' and epoch_acc > best_acc:                best_acc = epoch_acc                best_model_wts = 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))    # load best model weights    model.load_state_dict(best_model_wts)    return model

Visualizing the model predictions

def visualize_model(model, num_images=6):    images_so_far = 0    fig = plt.figure()    for i, data in enumerate(dataloders['val']):        inputs, labels = data        if use_gpu:            inputs, labels = Variable(inputs.cuda()), Variable(labels.cuda())        else:            inputs, labels = Variable(inputs), Variable(labels)        outputs = model(inputs)        _, preds = torch.max(outputs.data, 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:                return