Pytorch_finetune代码解读

最近跑了一下王晋东博士迁移学习简明手册上的深度网络的 finetune 代码实现，在这里做一下笔记。
来源：Githup开源链接

总结代码的大体框架如下：
1.数据集选择：office31-------load data
2.模型选择：Alexnet、Resnet50-------load model
3.优化器设置
4.模型训练
5.模型检测

下面来一个模块一个模块分析：

1.load data

# -*- coding: utf-8 -*-

from torchvision import datasets, transforms
import torch
import os


def load_data(root_path, dir, batch_size, phase):
    transform_dict = {
        'src': transforms.Compose(
        [transforms.RandomResizedCrop(224),#先随机采集，然后对裁剪得到的图像缩放为同一大小
         transforms.RandomHorizontalFlip(),#以给定的概率随机水平旋转给定的PIL的图像，默认为0.5
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406],#归一化
                              std=[0.229, 0.224, 0.225]),
         ]),
        'tar': transforms.Compose(
        [transforms.Resize(224),
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.485, 0.456, 0.406],
                              std=[0.229, 0.224, 0.225]),
         ])}
    data = datasets.ImageFolder(root=os.path.join(root_path, dir), transform=transform_dict[phase])
    #数据加载器，结合了数据集和取样器，并且可以提供多个线程处理数据集。
    #在训练模型时使用到此函数，用来把训练数据分成多个小组，此函数每次抛出一组数据。直至把所有的数据都抛出。就是做一个数据的初始化。
    data_loader = torch.utils.data.DataLoader(data, batch_size=batch_size, shuffle=True, drop_last=False, num_workers=4)
    return data_loader

def load_train(root_path, dir, batch_size, phase):
    transform_dict = {
        'src': transforms.Compose(
            [transforms.RandomResizedCrop(224),
             transforms.RandomHorizontalFlip(),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                  std=[0.229, 0.224, 0.225]),
             ]),
        'tar': transforms.Compose(
            [transforms.Resize(224),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                  std=[0.229, 0.224, 0.225]),
             ])}
    data = datasets.ImageFolder(root=os.path.join(root_path, dir), transform=transform_dict[phase])
    train_size = int(0.8 * len(data))
    test_size = len(data) - train_size
    data_train, data_val = torch.utils.data.random_split(data, [train_size, test_size])
    train_loader = torch.utils.data.DataLoader(data_train, batch_size=batch_size, shuffle=True, drop_last=False, num_workers=4)
    val_loader = torch.utils.data.DataLoader(data_val, batch_size=batch_size, shuffle=True, drop_last=False, num_workers=4)
    return train_loader, val_loader

分析：
这块代码保存为data_loader.py文件，主要实现数据集的下载、分类、格式调整等功能。这部分代码比较固定，在做其他深度网络的finetune时，一般情况下可以直接copy。
torch.utils.data.DataLoader（）函数用来加载数据，起到承上起下的作用。把得到的数据集直接送入模型中。

2.load model

def load_model(name='alexnet'):
    if name == 'alexnet':
        model = torchvision.models.alexnet(pretrained=True)
        n_features = model.classifier[6].in_features
        #定义预训练网络倒数第一层参数为输入，全连接层输出10个类别
        fc = torch.nn.Linear(n_features, args.n_class)
        model.classifier[6] = fc
    elif name == 'resnet':
        model = torchvision.models.resnet50(pretrained=True)
        n_features = model.fc.in_features
        fc = torch.nn.Linear(n_features, args.n_class)
        model.fc = fc
    model.fc.weight.data.normal_(0, 0.005)
    model.fc.bias.data.fill_(0.1)
    return model

分析：
这里只提供了两种网络的选择，确定使用的网络之后，直接在Pytorch自带的库中下载已经训练好的网络参数即可。注意，一般情况要冻结前n层，这里我们的finetune实验比较简单，直接冻结全连接层以前的所有层即可，然后定义一个新的全连接层，并把冻结模型的最后一层参数作为全连接层的输入参数，输出定义为自己实验数据集的类别即可。
注意，不同网络的定义会有细微区别，这部分代码不能生搬硬套。

3.优化器设置

def get_optimizer(model_name):
    learning_rate = args.lr
    if model_name == 'alexnet':
        param_group = [
            {'params': model.features.parameters(), 'lr': learning_rate}]
        for i in range(6):
            param_group += [{'params': model.classifier[i].parameters(),
                             'lr': learning_rate}]
        #最后一层提高学习率
        param_group += [{'params': model.classifier[6].parameters(),
                         'lr': learning_rate * 10}]
    elif model_name == 'resnet':
        param_group = []
        for k, v in model.named_parameters():
            #用于判断键是否存在于字典中，如果键在字典dict里返回true，否则返回false
            if not k.__contains__('fc'):
                param_group += [{'params': v, 'lr': learning_rate}]
            else:
                param_group += [{'params': v, 'lr': learning_rate * 10}]

    optimizer = optim.SGD(param_group, momentum=args.momentum)
    return optimizer

分析：
设置每一层参数的优化，注意最后一层的学习率与冻结层的学习率不同。

4.模型训练

def finetune(model, dataloaders, optimizer):
    since = time.time()
    best_acc = 0
    criterion = nn.CrossEntropyLoss()
    stop = 0

    for epoch in range(1, args.n_epoch + 1):
        stop += 1
        # You can uncomment this line for scheduling learning rate
        # lr_schedule(optimizer, epoch)
        #循环
        #训练集，测试集，验证集
        for phase in ['src', 'val', 'tar']:
            if phase == 'src':
                model.train()
            else:
                model.eval()
            total_loss, correct = 0, 0

            #一批一批算 batch_size
            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
                # 把梯度置零，也就是把loss关于weight的导数变成0
                optimizer.zero_grad()
                with torch.set_grad_enabled(phase == 'src'):
                    outputs = model(inputs)
                    loss = criterion(outputs, labels)
                preds = torch.max(outputs, 1)[1]
                if phase == 'src':
                    loss.backward()
                    # 更新所有的参数，一旦梯度被如backward()之类的函数计算好后，我们就可以调用这个函数
                    optimizer.step()
                total_loss += loss.item() * inputs.size(0)
                correct += torch.sum(preds == labels.data)
            #存储一次训练的平均loss和acc
            epoch_loss = total_loss / len(dataloaders[phase].dataset)
            epoch_acc = correct.double() / len(dataloaders[phase].dataset)
            print('Epoch: [{:02d}/{:02d}]---{}, loss: {:.6f}, acc: {:.4f}'.format(epoch, args.n_epoch, phase, epoch_loss,
                                                                                  epoch_acc))
            if phase == 'val' and epoch_acc > best_acc:
                stop = 0
                best_acc = epoch_acc
                torch.save(model.state_dict(), 'model.pkl')
        #如果20个训练还没有提高准确率，则跳出循环
        if stop >= args.early_stop:
            break
        print()

    model.load_state_dict(torch.load('model.pkl'))
    #使用训练好的的模型测试
    acc_test = test(model, dataloaders['tar'])
    time_pass = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_pass // 60, time_pass % 60))
    return model, acc_test

分析：
注意这里参数为：model、dataloaders、optimizers，即上面三部分做的工作，在这块代码全部用到。
注意训练集与测试集、验证集的运行有细微差别，比如训练集由loss、backward部分，但测试集和验证集则没有。

下面附上finetune_office31,py的全部代码，配合第一块data_loader.py使用

# -*- coding: utf-8 -*-

from __future__ import print_function

import argparse

import data_loader
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import time


# Command setting
parser = argparse.ArgumentParser(description='Finetune')
parser.add_argument('--model', type=str, default='resnet')
parser.add_argument('--batchsize', type=int, default=64)
parser.add_argument('--src', type=str, default='amazon')
parser.add_argument('--tar', type=str, default='webcam')
parser.add_argument('--n_class', type=int, default=31)
parser.add_argument('--lr', type=float, default=1e-4)
parser.add_argument('--n_epoch', type=int, default=100)
parser.add_argument('--momentum', type=float, default=0.9)
parser.add_argument('--decay', type=float, default=5e-4)
parser.add_argument('--data', type=str, default='Your dataset folder')
parser.add_argument('--early_stop', type=int, default=20)
args = parser.parse_args()

# Parameter setting
DEVICE = torch.device('cuda:2')
BATCH_SIZE = {'src': int(args.batchsize), 'tar': int(args.batchsize)}

#选择模型
def load_model(name='alexnet'):
    if name == 'alexnet':
        model = torchvision.models.alexnet(pretrained=True)
        n_features = model.classifier[6].in_features
        #定义预训练网络倒数第一层参数为输入，全连接层输出10个类别
        fc = torch.nn.Linear(n_features, args.n_class)
        model.classifier[6] = fc
    elif name == 'resnet':
        model = torchvision.models.resnet50(pretrained=True)
        n_features = model.fc.in_features
        fc = torch.nn.Linear(n_features, args.n_class)
        model.fc = fc
    model.fc.weight.data.normal_(0, 0.005)
    model.fc.bias.data.fill_(0.1)
    return model

#获得优化器
def get_optimizer(model_name):
    learning_rate = args.lr
    if model_name == 'alexnet':
        param_group = [
            {'params': model.features.parameters(), 'lr': learning_rate}]
        for i in range(6):
            param_group += [{'params': model.classifier[i].parameters(),
                             'lr': learning_rate}]
        #最后一层提高学习率
        param_group += [{'params': model.classifier[6].parameters(),
                         'lr': learning_rate * 10}]
    elif model_name == 'resnet':
        param_group = []
        for k, v in model.named_parameters():
            #用于判断键是否存在于字典中，如果键在字典dict里返回true，否则返回false
            if not k.__contains__('fc'):
                param_group += [{'params': v, 'lr': learning_rate}]
            else:
                param_group += [{'params': v, 'lr': learning_rate * 10}]

    optimizer = optim.SGD(param_group, momentum=args.momentum)
    return optimizer


# Schedule learning rate
#运行时可以先不管这部分，后续优化时会用到
def lr_schedule(optimizer, epoch):
    def lr_decay(LR, n_epoch, e):
        return LR / (1 + 10 * e / n_epoch) ** 0.75
    #调整学习率
    for i in range(len(optimizer.param_groups)):
        if i < len(optimizer.param_groups) - 1:
            optimizer.param_groups[i]['lr'] = lr_decay(
                args.lr, args.n_epoch, epoch)
        else:
            optimizer.param_groups[i]['lr'] = lr_decay(
                args.lr, args.n_epoch, epoch) * 10

def test(model, target_test_loader):
    model.eval()
    correct = 0
    criterion = torch.nn.CrossEntropyLoss()
    len_target_dataset = len(target_test_loader.dataset)

    #在使用pytorch时，并不是所有的操作都需要进行计算图的生成（计算过程的构建，以便梯度反向传播等操作）。
    #而对于tensor的计算操作，默认是要进行计算图的构建的，
    #在这种情况下，可以使用 with torch.no_grad():，强制之后的内容不进行计算图构建。

    with torch.no_grad():
        for data, target in target_test_loader:
            data, target = data.to(DEVICE), target.to(DEVICE)
            s_output = model(data)
            loss = criterion(s_output, target)
            pred = torch.max(s_output, 1)[1]
            correct += torch.sum(pred == target)
    acc = correct.double() / len(target_test_loader.dataset)
    return acc


"""
            #torch.max(input, dim) 函数
            输入：
            input是softmax函数输出的一个tensor
            dim是max函数索引的维度0 / 1，0
            是每列的最大值，1
            是每行的最大值
            输出：
            函数会返回两个tensor，
            第一个tensor是每行的最大值，softmax的输出中最大的是1，
            所以第一个tensor是全1的tensor；第二个tensor是每行最大值的索引。

"""

def finetune(model, dataloaders, optimizer):
    since = time.time()
    best_acc = 0
    criterion = nn.CrossEntropyLoss()
    stop = 0

    for epoch in range(1, args.n_epoch + 1):
        stop += 1
        # You can uncomment this line for scheduling learning rate
        # lr_schedule(optimizer, epoch)
        #循环
        #训练集，测试集，验证集
        for phase in ['src', 'val', 'tar']:
            if phase == 'src':
                model.train()
            else:
                model.eval()
            total_loss, correct = 0, 0

            #一批一批算 batch_size
            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
                # 把梯度置零，也就是把loss关于weight的导数变成0
                optimizer.zero_grad()
                with torch.set_grad_enabled(phase == 'src'):
                    outputs = model(inputs)
                    loss = criterion(outputs, labels)
                preds = torch.max(outputs, 1)[1]
                if phase == 'src':
                    loss.backward()
                    # 更新所有的参数，一旦梯度被如backward()之类的函数计算好后，我们就可以调用这个函数
                    optimizer.step()
                total_loss += loss.item() * inputs.size(0)
                correct += torch.sum(preds == labels.data)
            #存储一次训练的平均loss和acc
            epoch_loss = total_loss / len(dataloaders[phase].dataset)
            epoch_acc = correct.double() / len(dataloaders[phase].dataset)
            print('Epoch: [{:02d}/{:02d}]---{}, loss: {:.6f}, acc: {:.4f}'.format(epoch, args.n_epoch, phase, epoch_loss,
                                                                                  epoch_acc))
            if phase == 'val' and epoch_acc > best_acc:
                stop = 0
                best_acc = epoch_acc
                torch.save(model.state_dict(), 'model.pkl')
        #如果20个训练还没有提高准确率，则跳出循环
        if stop >= args.early_stop:
            break
        print()

    model.load_state_dict(torch.load('model.pkl'))
    #使用训练好的的模型测试
    acc_test = test(model, dataloaders['tar'])
    time_pass = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_pass // 60, time_pass % 60))
    return model, acc_test


if __name__ == '__main__':
    torch.manual_seed(10)
    # Load data
    root_dir = args.data
    domain = {'src': str(args.src), 'tar': str(args.tar)}
    dataloaders = {}
    dataloaders['tar'] = data_loader.load_data(root_dir, domain['tar'], BATCH_SIZE['tar'], 'tar')
    dataloaders['src'], dataloaders['val'] = data_loader.load_train(root_dir, domain['src'], BATCH_SIZE['src'], 'src')
    # Load model
    model_name = str(args.model)
    #device=
    model = load_model(model_name).to(DEVICE)
    print('Source: {} ({}), target: {} ({}), model: {}'.format(
        domain['src'], len(dataloaders['src'].dataset), domain['tar'], len(dataloaders['val'].dataset), model_name))

    optimizer = get_optimizer(model_name)
    model_best, best_acc = finetune(model, dataloaders, optimizer)
    print('Best acc: {}'.format(best_acc))

打印每次训练的准确率：

end