修改pytorch官方实例适用于自己的二分类迁移学习项目

本demo从pytorch官方的迁移学习示例修改而来，增加了以下功能：

1. 根据AUC来迭代最优参数；
2. 五折交叉验证；
3. 输出验证集错误分类图片；
4. 输出分类报告并保存AUC结果图片。

    import os
    import numpy as np
    import torch
    import torch.nn as nn
    from torch.optim import lr_scheduler
    import torchvision
    from torchvision import datasets, models, transforms
    from torch.utils.data import DataLoader
    from sklearn.metrics import roc_auc_score, classification_report
    from sklearn.model_selection import KFold
    from torch.autograd import Variable
    import torch.optim as optim
    import time
    import copy
    import shutil
    import sys
    import scikitplot as skplt
    import matplotlib.pyplot as plt
    import pandas as pd
   
    plt.switch_backend('agg')
    N_CLASSES = 2
    BATCH_SIZE = 8
    DATA_DIR = './data'
    LABEL_DICT = {0: 'class_1', 1: 'class_2'}
   
    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(100)
   
    def train_model(model, criterion, optimizer, scheduler, fold, name, num_epochs=25):
        since = time.time()
        # 先深拷贝一份当前模型的参数，后面迭代过程中若遇到更优模型则替换
        best_model_wts = copy.deepcopy(model.state_dict())
        # best_acc = 0.0
        # 初始auc
        best_auc = 0.0
        best_desc = [0, 0, None]
        best_img_name = None
        plt_auc = [None, None]
   
        for epoch in range(num_epochs):
            print('Epoch {}/{}'.format(epoch, num_epochs - 1))
            print('- ' * 50)
   
            for phase in ['train', 'val']:
                if phase == 'train':
                    # 训练的时候进行学习率规划，其定义在下面给出
                    scheduler.step()
                    model.train(True)
                else:
                    model.train(False)
                phase_pred = np.array([])
                phase_label = np.array([])
                img_name = np.zeros((1, 2))
                prob_pred = np.zeros((1, 2))
                running_loss = 0.0
                running_corrects = 0
                # 这样迭代方便跟踪图片路径，输出错误图片名称
                for data, index in zip(dataloaders[phase], dataloaders[phase].batch_sampler):
                    inputs, labels = data
                    if use_gpu:
                        inputs = Variable(inputs.cuda())
                        labels = Variable(labels.cuda())
                    else:
                        inputs, labels = Variable(inputs), Variable(labels)
   
                    # 梯度参数设为0
                    optimizer.zero_grad()
   
                    # forward
                    outputs = model(inputs)
                    _, preds = torch.max(outputs.data, 1)
                    loss = criterion(outputs, labels)
   
                    # backward + 训练阶段优化
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()
   
                    if phase == 'val':
                        img_name = np.append(img_name, np.array(dataloaders[phase].dataset.imgs)[index], axis=0)
                        prob = outputs.data.cpu().numpy()
                        prob_pred = np.append(prob_pred, prob, axis=0)
   
                    phase_pred = np.append(phase_pred, preds.cpu().numpy())
                    phase_label = np.append(phase_label, labels.data.cpu().numpy())
                    running_loss += loss.item() * inputs.size(0)
                    running_corrects += torch.sum(preds == labels.data).float()
                print()
                epoch_loss = running_loss / dataset_sizes[phase]
                epoch_acc = running_corrects / dataset_sizes[phase]
                epoch_auc = roc_auc_score(phase_label, phase_pred)
                print('{} Loss: {:.4f} Acc: {:.4f} Auc: {:.4f}'.format(
                    phase, epoch_loss, epoch_acc, epoch_auc))
                report = classification_report(phase_label, phase_pred, target_names=class_names)
                print(report)
   
                img_name = zip(img_name[1:], phase_pred)
                # 当验证时遇到了更好的模型则予以保留
                if phase == 'val' and epoch_auc > best_auc:
                    best_auc = epoch_auc
                    best_desc = epoch_acc, epoch_auc, report
                    best_img_name = img_name
                    # 深拷贝模型参数
                    best_model_wts = copy.deepcopy(model.state_dict())
                    plt_auc = phase_label, prob_pred[1:]
   
            print()
        print(plt_auc[0].shape, plt_auc[1].shape)
        csv_file = pd.DataFrame(plt_auc[1], columns=['class_1', 'class_2'])
        csv_file['true_label'] = pd.DataFrame(plt_auc[0])
        csv_file['true_label'] = csv_file['true_label'].apply(lambda x: LABEL_DICT[x])
        csv_file.to_csv(f'./prob_result/{name}_fold_{fold}_porb.csv', index=False)
        skplt.metrics.plot_roc_curve(plt_auc[0], plt_auc[1], curves=['each_class'])
        plt.savefig(f'./roc_img/{name}_fold_{fold}_roc.png', dpi=600)
        time_elapsed = time.time() - since
        print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
        reports = 'The Desc according to the Best val Auc: \nACC -> {:4f}\nAclass_2 -> {:4f}\n\n{}'.format(best_desc[0], best_desc[1],
                                                                                             best_desc[2])
        report_file.write(reports)
        print(reports)
        print('List the wrong judgement img ...')
        count = 0
        for i in best_img_name:
            actual_label = int(i[0][1])
            pred_label = i[1]
            if actual_label != pred_label:
                tmp_word = f'{i[0][0].split("/")[-1]}, actual: {LABEL_DICT[actual_label]}, ' \
                           f'pred: {LABEL_DICT[pred_label]}'
                print(tmp_word)
                label_file.write(tmp_word + '\n')
                count += 1
        print(f'This fold has {count} wrong records ...')
   
        # 载入最优模型参数
        model.load_state_dict(best_model_wts)
        return model
   
    def plot_img():
        for i, data in enumerate(dataloaders['train']):
            inputs, classes = data
            out = torchvision.utils.make_grid(inputs)
            imshow(out, title=[class_names[x] for x in classes])
   
    # 此函数可以修改适用于自己项目的图片文件名
    def move_file(data, file_path, dir_path, root_path):
        label_0 = 'class_2'
        label_1 = 'class_1'
        print(f'start copy the {file_path} file ...')
        os.chdir(dir_path)
        if os.path.exists(file_path):
            print(f'Find exist {file_path} file, the file will be dropped.')
            shutil.rmtree(os.path.join(root_path, dir_path, file_path))
            print(f'Finish drop the {file_path} file.')
   
        os.mkdir(file_path)
        tmp_path = os.path.join(os.getcwd(), file_path)
        tmp_pre_path = os.getcwd()
        for d in data:
            pre_path = os.path.join(tmp_pre_path, d)
            os.chdir(tmp_path)
            if d[:2] == label_0:
                if not os.path.exists(label_0):
                    os.mkdir(label_0)
                cur_path = os.path.join(tmp_path, label_0, d)
                shutil.copyfile(pre_path, cur_path)
            if d[:2] == label_1:
                if not os.path.exists(label_1):
                    os.mkdir(label_1)
                cur_path = os.path.join(tmp_path, label_1, d)
                shutil.copyfile(pre_path, cur_path)
        print('finish this work ...')
   
    if __name__ == "__main__":
        if not os.path.exists('roc_img'):
            os.mkdir('roc_img')
        if not os.path.exists('prob_result'):
            os.mkdir('prob_result')
        if not os.path.exists('report'):
            os.mkdir('report')
        if not os.path.exists('error_record'):
            os.mkdir('error_record')
        if not os.path.exists('model'):
            os.mkdir('model')
        label_file = open(f'./error_record/{sys.argv[1]}_img_name_actual_pred.txt', 'w')
   
        kf = KFold(n_splits=5, shuffle=True, random_state=1)
        origin_path = '/home/project/'
        dd_list = np.array([o for o in os.listdir(DATA_DIR) if os.path.isfile(os.path.join(DATA_DIR, o))])
   
        for m, n in enumerate(kf.split(dd_list), start=1):
            report_file = open(f'./report/{sys.argv[1]}_fold_{m}_report.txt', 'w')
            print(f'The {m} fold for copy file and training ...')
            move_file(dd_list[n[0]], 'train', DATA_DIR, origin_path)
            os.chdir(origin_path)
            move_file(dd_list[n[1]], 'val', DATA_DIR, origin_path)
            os.chdir(origin_path)
            data_transforms = {
                'train': transforms.Compose([
                    # 裁剪到224,224
                    transforms.RandomResizedCrop(224),
                    # 随机水平翻转给定的PIL.Image,概率为0.5。即：一半的概率翻转，一半的概率不翻转。
                    transforms.RandomHorizontalFlip(),
                    # transforms.ColorJitter(0.05, 0.05, 0.05, 0.05),  # HSV以及对比度变化
                    transforms.ToTensor(),
                    # 把一个取值范围是[0,255]的PIL.Image或者shape为(H,W,C)的numpy.ndarray，转换成形状为[C,H,W]，取值范围是[0,1.0]的FloadTensor
                    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])
                ]),
            }
   
            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=BATCH_SIZE,
                                                          shuffle=True, num_workers=8, pin_memory=False)
                           for x in ['train', 'val']}
   
            dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
   
            class_names = image_datasets['train'].classes
            size = len(class_names)
            print('label mapping: ')
            print(image_datasets['train'].class_to_idx)
            use_gpu = torch.cuda.is_available()
            model_ft = None
            if sys.argv[1] == 'resnet':
                model_ft = models.resnet50(pretrained=True)
                num_ftrs = model_ft.fc.in_features
                model_ft.fc = nn.Sequential(
                    nn.Linear(num_ftrs, N_CLASSES),
                    nn.Sigmoid()
                )
   
            # 这边可以自行把inception模型加进去
            if sys.argv[1] == 'inception':
                raise Exception("not provide inception model ...")
                # model_ft = models.inception_v3(pretrained=True)
   
            if sys.argv[1] == 'desnet':
                model_ft = models.densenet121(pretrained=True)
                num_ftrs = model_ft.classifier.in_features
                model_ft.classifier = nn.Sequential(
                    nn.Linear(num_ftrs, N_CLASSES),
                    nn.Sigmoid()
                )
                # use_gpu = False
   
            if use_gpu:
                model_ft = model_ft.cuda()
   
            criterion = nn.CrossEntropyLoss()
            optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
            # 每7个epoch衰减0.1倍
            exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
            model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, m, sys.argv[1], num_epochs=25)
            print('Start save the model ...')
            torch.save(model_ft.state_dict(), f'./model/fold_{m}_{sys.argv[1]}.pkl')
            print(f'The mission of the fold {m} finished.')
            print('# '*50)
            report_file.close()
        label_file.close()