0025-pytorch入门-分类模型经典案例

缺乏、安全感 2023-01-09 12:34 220阅读 0赞

### 参数设置，config.py ###

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/20 19:07
    # @Author  : Johnson
    class DefaultConfigs(object):
        #1.string parameters
        train_data = "/home/user/zls/data/train/"
        test_data = ""
        val_data = "/home/user/zls/data/val/"
        model_name = "resnet50"
        weights = "./checkpoints/"
        best_models = weights + "best_model/"
        submit = "./submit/"
        logs = "./logs/"
        gpus = "1"
        augmen_level = "medium"  # "light","hard","hard2"
    
        #2.numeric parameters
        epochs = 40
        batch_size = 16
        img_height = 300
        img_weight = 300
        num_classes = 62
        seed = 888
        lr = 1e-4
        lr_decay = 1e-4
        weight_decay = 1e-4
    
    config = DefaultConfigs()

### 数据加载：data\_loader.py ###

pytorch 的数据读取方式有两种，一种是不同类别的图像按照文件夹进行划分，比如交通标志数据集：

train/

00000/

01153\_00000.png  
01153\_00001.png  
00001/

00025\_00000.png  
00025\_00001.png

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/20 19:01
    # @Author  : Johnson
    from torch.utils.data import Dataset
    from torchvision import transforms as T
    from config import config
    from PIL import Image
    from itertools import chain
    from glob import glob
    from tqdm import tqdm
    from .augmentations import get_train_transform,get_test_transform
    import random
    import numpy as np
    import pandas as pd
    import os
    import cv2
    import torch
    
    #1.set random seed
    random.seed(config.seed)
    np.random.seed(config.seed)
    torch.manual_seed(config.seed)
    torch.cuda.manual_seed_all(config.seed)
    
    #2.define dataset
    class ChaojieDataset(Dataset):
        def __init__(self,label_list,train=True,test=False):
            self.test = test
            self.train = train
            imgs = []
            if self.test:
                for index,row in label_list.iterrows():
                    imgs.append(row["filename"])
                self.imgs = imgs
            else:
                for index,row in label_list.iterrows():
                    imgs.append(row['filename',row["label"]])
                self.imgs = imgs
    
    
        def __getitem__(self,index):
            if self.test:
                filename = self.imgs[index]
                img = cv2.imread(filename)
                img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
                img = cv2.resize(img,(int(config.img_height*1.5),int(config.img_weight*1.5)))
                img = get_test_transform(img.shape)(image=img)["image"]
                return img,filename
            else:
                filename,label = self.imgs[index]
                img = cv2.imread(filename)
                img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
                img = cv2.resize(img, (int(config.img_height * 1.5), int(config.img_weight * 1.5)))
                img = get_train_transform(img.shape, augmentation=config.augmen_level)(image=img)["image"]
                return img, label
        def __len__(self):
            return len(self.imgs)
    
    
    def collate_fn(batch):
        imgs = []
        label = []
        for sample in batch:
            imgs.append(sample[0])
            label.append(sample[1])
    
        return torch.stack(imgs,0),label
    
    
    def get_files(root,mode):
        #for test
        if mode == "test":
            files = []
            for img in os.listdir(root):
                files.append(root + img)
            files = pd.DataFrame({"filename":files})
            return files
        elif mode != "test":
            #for train and val
            all_data_path,labels = [],[]
            image_folders = list(map(lambda x:root+x,os.listdir(root)))
            all_images = list(chain.from_iterable(list(map(lambda x:glob(x+"/*"),image_folders))))
            print("loading train dataset")
            for file in tqdm(all_images):
                all_data_path.append(file)
                labels.append(int(file.split("/")[-2]))
            all_files = pd.DataFrame({"filename":all_data_path,"label":labels})
            return all_files
        else:
            print("check the mode please!")

#### 数据增强：augmentations.py ####

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/20 19:00
    # @Author  : Johnson
    import random
    from typing import List,Tuple
    
    import albumentations as A
    import cv2
    import numpy as np
    from albumentations.augmentations.functional import brightness_contrast_adjust,elastic_transform
    from albumentations.pytorch import ToTensor
    from albumentations.augmentations.transforms import Resize,CenterCrop
    from config import config
    
    class IndependentRandomBrightnessContrast(A.ImageOnlyTransform):
        """ Change brightness & contrast independently per channels """
    
        def __init__(self, brightness_limit=0.2, contrast_limit=0.2, always_apply=False, p=0.5):
            super(IndependentRandomBrightnessContrast, self).__init__(always_apply, p)
            self.brightness_limit = A.to_tuple(brightness_limit)
            self.contrast_limit = A.to_tuple(contrast_limit)
    
        def apply(self, img, **params):
            img = img.copy()
            for ch in range(img.shape[2]):
                alpha = 1.0 + random.uniform(self.contrast_limit[0], self.contrast_limit[1])
                beta = 0.0 + random.uniform(self.brightness_limit[0], self.brightness_limit[1])
                img[..., ch] = brightness_contrast_adjust(img[..., ch], alpha, beta)
    
            return img
    
    def get_light_augmentations(image_size):
        return A.Compose([
            A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1,
                               rotate_limit=15,
                               border_mode=cv2.BORDER_CONSTANT, value=0),
            A.RandomSizedCrop(min_max_height=(int(image_size[0] * 0.85), image_size[0]),
                              height=image_size[0],
                              width=image_size[1], p=0.3),
            # Brightness/contrast augmentations
            A.OneOf([
                A.RandomBrightnessContrast(brightness_limit=0.25,
                                           contrast_limit=0.2),
                IndependentRandomBrightnessContrast(brightness_limit=0.1,
                                                    contrast_limit=0.1),
                A.RandomGamma(gamma_limit=(75, 125)),
                A.NoOp()
            ]),
            A.HorizontalFlip(p=0.5),
        ])
    
    def get_medium_augmentations(image_size):
        return A.Compose([
            A.OneOf([
                A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1,
                                   rotate_limit=15,
                                   border_mode=cv2.BORDER_CONSTANT, value=0),
                A.OpticalDistortion(distort_limit=0.11, shift_limit=0.15,
                                    border_mode=cv2.BORDER_CONSTANT,
                                    value=0),
                A.NoOp()
            ]),
            A.RandomSizedCrop(min_max_height=(int(image_size[0] * 0.75), image_size[0]),
                              height=image_size[0],
                              width=image_size[1], p=0.3),
            A.OneOf([
                A.RandomBrightnessContrast(brightness_limit=0.5,
                                           contrast_limit=0.4),
                IndependentRandomBrightnessContrast(brightness_limit=0.25,
                                                    contrast_limit=0.24),
                A.RandomGamma(gamma_limit=(50, 150)),
                A.NoOp()
            ]),
            A.OneOf([
                A.RGBShift(r_shift_limit=20, b_shift_limit=15, g_shift_limit=15),
                A.HueSaturationValue(hue_shift_limit=5,
                                     sat_shift_limit=5),
                A.NoOp()
            ]),
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5)
        ])
    
    
    def get_hard_augmentations(image_size):
        return A.Compose([
            A.OneOf([
                A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1,
                                   rotate_limit=45,
                                   border_mode=cv2.BORDER_CONSTANT, value=0),
                A.ElasticTransform(alpha_affine=0,
                                   alpha=35,
                                   sigma=5,
                                   border_mode=cv2.BORDER_CONSTANT,
                                   value=0),
                A.OpticalDistortion(distort_limit=0.11, shift_limit=0.15,
                                    border_mode=cv2.BORDER_CONSTANT,
                                    value=0),
                A.GridDistortion(border_mode=cv2.BORDER_CONSTANT,
                                 value=0),
                A.NoOp()
            ]),
    
            A.OneOf([
    
                A.RandomSizedCrop(min_max_height=(int(image_size[0] * 0.75), image_size[0]),
                                  height=image_size[0],
                                  width=image_size[1], p=0.3),
                A.NoOp()
            ]),
    
            A.ISONoise(p=0.5),
    
            # Brightness/contrast augmentations
            A.OneOf([
                A.RandomBrightnessContrast(brightness_limit=0.5,
                                           contrast_limit=0.4),
                IndependentRandomBrightnessContrast(brightness_limit=0.25,
                                                    contrast_limit=0.24),
                A.RandomGamma(gamma_limit=(50, 150)),
                A.NoOp()
            ]),
    
            A.OneOf([
                A.RGBShift(r_shift_limit=40, b_shift_limit=30, g_shift_limit=30),
                A.HueSaturationValue(hue_shift_limit=10,
                                     sat_shift_limit=10),
                A.ToGray(p=0.2),
                A.NoOp()
            ]),
    
            A.ChannelDropout(),
            A.RandomGridShuffle(p=0.3),
    
            # D4
            A.Compose([
                A.RandomRotate90(),
                A.Transpose()
            ])
        ])
    
    def get_hard_augmentations_v2(image_size):
        return A.Compose([
            A.OneOf([
                A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.1,
                                   rotate_limit=45,
                                   border_mode=cv2.BORDER_CONSTANT, value=0),
                A.ElasticTransform(alpha_affine=0,
                                   alpha=35,
                                   sigma=5,
                                   border_mode=cv2.BORDER_CONSTANT,
                                   value=0),
                A.OpticalDistortion(distort_limit=0.11, shift_limit=0.15,
                                    border_mode=cv2.BORDER_CONSTANT,
                                    value=0),
                A.GridDistortion(border_mode=cv2.BORDER_CONSTANT,
                                 value=0),
                A.NoOp()
            ]),
    
            A.OneOf([
                A.RandomSizedCrop(min_max_height=(int(image_size[0] * 0.75), image_size[0]),
                                  height=image_size[0],
                                  width=image_size[1], p=0.3),
                A.NoOp()
            ]),
    
            A.ISONoise(p=0.5),
            A.JpegCompression(p=0.3, quality_lower=75),
    
            # Brightness/contrast augmentations
            A.OneOf([
                A.RandomBrightnessContrast(brightness_limit=0.5,
                                           contrast_limit=0.4),
                IndependentRandomBrightnessContrast(brightness_limit=0.25,
                                                    contrast_limit=0.24),
                A.RandomGamma(gamma_limit=(50, 150)),
                A.NoOp()
            ]),
    
            A.OneOf([
                A.RGBShift(r_shift_limit=40, b_shift_limit=30, g_shift_limit=30),
                A.HueSaturationValue(hue_shift_limit=10,
                                     sat_shift_limit=10),
                A.ToGray(p=0.2),
                A.NoOp()
            ]),
    
    
            A.OneOf([
                A.ChannelDropout(p=0.2),
                A.CoarseDropout(p=0.1, max_holes=2, max_width=256, max_height=256, min_height=16, min_width=16),
                A.NoOp()
            ]),
    
            A.RandomGridShuffle(p=0.3),
    
            # D4
            A.Compose([
                A.RandomRotate90(),
                A.Transpose()
            ])
        ])
    
    def get_none_augmentations(image_size):
        return A.NoOp()
    
    
    def get_train_transform(image_size,augmentation=None):
        if augmentation is None:
            augmentation = 'none'
    
        LEVELS = {
            'none':get_none_augmentations,
            'light':get_light_augmentations,
            'medium':get_medium_augmentations,
            'hard':get_hard_augmentations,
            'hard2':get_hard_augmentations_v2
        }
        assert augmentation is LEVELS.keys()
        augmentation = LEVELS[augmentation](image_size)
    
        longest_size = max(image_size[0],image_size[1])
        return A.Compose([
            # Resize(int(config.img_height*1.5),int(config.img_weight*1.5)),
            CenterCrop(config.img_height, config.img_weight),
            A.LongestMaxSize(longest_size, interpolation=cv2.INTER_CUBIC),
    
            A.PadIfNeeded(image_size[0], image_size[1],
                          border_mode=cv2.BORDER_CONSTANT, value=0),
            augmentation,
            A.Normalize(),
            ToTensor()
        ])
    
    def get_test_transform(image_size):
        longest_size = max(image_size[0], image_size[1])
        return A.Compose([
            #Resize(int(config.img_height*1.5),int(config.img_weight*1.5)),
            CenterCrop(config.img_height,config.img_weight),
            A.LongestMaxSize(longest_size, interpolation=cv2.INTER_CUBIC),
    
            A.PadIfNeeded(image_size[0], image_size[1],
                          border_mode=cv2.BORDER_CONSTANT, value=0),
    
            A.Normalize(),
            ToTensor()
        ])

### 数据加载:model.py ###

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/20 20:55
    # @Author  : Johnson
    import torchvision
    import torch.nn.functional as F
    from torch import nn
    from config import config
    
    def generate_model():
        class DenseModel(nn.Module):
            def __init__(self, pretrained_model):
                super(DenseModel, self).__init__()
                self.classifier = nn.Linear(pretrained_model.classifier.in_features, config.num_classes)
    
                for m in self.modules():
                    if isinstance(m, nn.Conv2d):
                        nn.init.kaiming_normal(m.weight)
                    elif isinstance(m, nn.BatchNorm2d):
                        m.weight.data.fill_(1)
                        m.bias.data.zero_()
                    elif isinstance(m, nn.Linear):
                        m.bias.data.zero_()
    
                self.features = pretrained_model.features
                self.layer1 = pretrained_model.features._modules['denseblock1']
                self.layer2 = pretrained_model.features._modules['denseblock2']
                self.layer3 = pretrained_model.features._modules['denseblock3']
                self.layer4 = pretrained_model.features._modules['denseblock4']
    
    
            def forward(self,x):
                features = self.features(x)
                out = F.relu(features,inplace=True)
                out = F.avg_pool2d(out,kernel_size=8).view(features.size(0),-1)
                out = F.sigmoid(self.classifier(out))
                return out
    
    
    
        return DenseModel(torchvision.models.densenet169(pretrained=True))
    
    def get_net():
        #return MyModel(torchvision.models.resnet101(pretrained = True))
        model = torchvision.models.resnet152(pretrained = True)
        #for param in model.parameters():
        #    param.requires_grad = False
        model.avgpool = nn.AdaptiveAvgPool2d(1)
        model.fc = nn.Linear(2048,config.num_classes)
        return model

### 指标评价：utils.py ###

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/20 21:04
    # @Author  : Johnson
    import shutil
    import torch
    import sys
    import os
    from config import config
    
    def save_checkpoint(state, is_best,fold):
        filename = config.weights + config.model_name + os.sep +str(fold) + os.sep + "_checkpoint.pth.tar"
        torch.save(state, filename)
        if is_best:
            message = config.best_models + config.model_name+ os.sep +str(fold)  + os.sep + 'model_best.pth.tar'
            print("Get Better top1 : %s saving weights to %s"%(state["best_precision1"],message))
            with open("./logs/%s.txt"%config.model_name,"a") as f:
                print("Get Better top1 : %s saving weights to %s"%(state["best_precision1"],message),file=f)
            shutil.copyfile(filename, message)
    
    
    
    class AverageMeter(object):
        """Computes and stores the average and current value"""
        def __init__(self):
            self.reset()
    
        def reset(self):
            self.val = 0
            self.avg = 0
            self.sum = 0
            self.count = 0
    
        def update(self, val, n=1):
            self.val = val
            self.sum += val * n
            self.count += n
            self.avg = self.sum / self.count
    
    
    def adjust_learning_rate(optimizer,epoch):
        """sets the learning rate to the initial LR decayed by 10 every 3 epochs """
        lr = config.lr*(0.1**(epoch//3))
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr
    
    def schedule(current_epoch, current_lrs, **logs):
        lrs = [1e-3, 1e-4, 0.5e-4, 1e-5, 0.5e-5]
        epochs = [0, 1, 6, 8, 12]
        for lr, epoch in zip(lrs, epochs):
            if current_epoch >= epoch:
                current_lrs[5] = lr
                if current_epoch >= 2:
                    current_lrs[4] = lr * 1
                    current_lrs[3] = lr * 1
                    current_lrs[2] = lr * 1
                    current_lrs[1] = lr * 1
                    current_lrs[0] = lr * 0.1
        return current_lrs
    
    
    def accuracy(output, target, topk=(1,)):
        """Computes the accuracy over the k top predictions for the specified values of k"""
        with torch.no_grad():
            maxk = max(topk)
            batch_size = target.size(0)
    
            _, pred = output.topk(maxk, 1, True, True)
            pred = pred.t()
            correct = pred.eq(target.view(1, -1).expand_as(pred))
    
            res = []
            for k in topk:
                correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
                res.append(correct_k.mul_(100.0 / batch_size))
            return res
    
    
    def get_learning_rate(optimizer):
        lr=[]
        for param_group in optimizer.param_groups:
           lr +=[ param_group['lr'] ]
    
        #assert(len(lr)==1) #we support only one param_group
        lr = lr[0]
    
        return lr
    
    
    def time_to_str(t, mode='min'):
        if mode=='min':
            t  = int(t)/60
            hr = t//60
            min = t%60
            return '%2d hr %02d min'%(hr,min)
    
        elif mode=='sec':
            t   = int(t)
            min = t//60
            sec = t%60
            return '%2d min %02d sec'%(min,sec)
    
    
        else:
            raise NotImplementedError

### 进度条:progress.py ###

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/21 9:59
    # @Author  : Johnson
    import sys
    import re
    
    
    class ProgressBar(object):
        DEFAULT = "Progress: %(bar)s %(percent)3d%%"
    
        def __init__(self, mode, epoch=None, total_epoch=None, current_loss=None, current_top1=None, model_name=None,
                     total=None, current=None, width=50, symbol=">", output=sys.stderr):
            assert len(symbol) == 1
    
            self.mode = mode
            self.total = total
            self.symbol = symbol
            self.output = output
            self.width = width
            self.current = current
            self.epoch = epoch
            self.total_epoch = total_epoch
            self.current_loss = current_loss
            self.current_top1 = current_top1
            self.model_name = model_name
    
        def __call__(self):
            percent = self.current / float(self.total)
            size = int(self.width * percent)
            bar = "[" + self.symbol * size + " " * (self.width - size) + "]"
    
            args = {
                "mode": self.mode,
                "total": self.total,
                "bar": bar,
                "current": self.current,
                "percent": percent * 100,
                "current_loss": self.current_loss,
                "current_top1": self.current_top1,
                "epoch": self.epoch + 1,
                "epochs": self.total_epoch
            }
            message = "\033[1;32;40m%(mode)s Epoch:  %(epoch)d/%(epochs)d %(bar)s\033[0m  [Current: Loss %(current_loss)f Top1: %(current_top1)f ]  %(current)d/%(total)d \033[1;32;40m[ %(percent)3d%% ]\033[0m" % args
            self.write_message = "%(mode)s Epoch:  %(epoch)d/%(epochs)d %(bar)s  [Current: Loss %(current_loss)f Top1: %(current_top1)f ]  %(current)d/%(total)d [ %(percent)3d%% ]" % args
            print("\r" + message, file=self.output, end="")
    
        def done(self):
            self.current = self.total
            self()
            print("", file=self.output)
            with open("./logs/%s.txt" % self.model_name, "a") as f:
                print(self.write_message, file=f)
    
    
    if __name__ == "__main__":
    
        from time import sleep
    
        progress = ProgressBar("Train", total_epoch=150, model_name="resnet159")
        for i in range(150):
            progress.total = 50
            progress.epoch = i
            progress.current_loss = 0.15
            progress.current_top1 = 0.45
            for x in range(50):
                progress.current = x
                progress()
                sleep(0.1)
            progress.done()

### 主函数:main.py ###

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/21 10:02
    # @Author  : Johnson
    import os
    import random
    import time
    import json
    import torch
    import torchvision
    import numpy as np
    import pandas as pd
    import warnings
    from datetime import datetime
    from torch import nn,optim
    from config import config
    from collections import OrderedDict
    from torch.autograd import Variable
    from torch.utils.data import DataLoader
    from dataset.dataloader import *
    from sklearn.model_selection import train_test_split,StratifiedKFold
    from timeit import default_timer as timer
    from models.model import *
    from utils import *
    from IPython import embed
    #1. set random.seed and cudnn performance
    random.seed(config.seed)
    np.random.seed(config.seed)
    torch.manual_seed(config.seed)
    torch.cuda.manual_seed_all(config.seed)
    os.environ["CUDA_VISIBLE_DEVICES"] = config.gpus
    torch.backends.cudnn.benchmark = True
    warnings.filterwarnings('ignore')
    
    
    # 2. evaluate func
    def evaluate(val_loader, model, criterion, epoch):
        # 2.1 define meters
        losses = AverageMeter()
        top1 = AverageMeter()
        # progress bar
        val_progressor = ProgressBar(mode="Val  ", epoch=epoch, total_epoch=config.epochs, model_name=config.model_name,
                                     total=len(val_loader))
        # 2.2 switch to evaluate mode and confirm model has been transfered to cuda
        model.cuda()
        model.eval()
        with torch.no_grad():
            for i, (input, target) in enumerate(val_loader):
                val_progressor.current = i
                input = Variable(input).cuda()
                target = Variable(torch.from_numpy(np.array(target)).long()).cuda()
                # target = Variable(target).cuda()
                # 2.2.1 compute output
                output = model(input)
                loss = criterion(output, target)
    
                # 2.2.2 measure accuracy and record loss
                precision1, precision2 = accuracy(output, target, topk=(1, 2))
                losses.update(loss.item(), input.size(0))
                top1.update(precision1[0], input.size(0))
                val_progressor.current_loss = losses.avg
                val_progressor.current_top1 = top1.avg
                val_progressor()
            val_progressor.done()
        return [losses.avg, top1.avg]
    
    
    # 3. test model on public dataset and save the probability matrix
    def test(test_loader, model, folds):
        # 3.1 confirm the model converted to cuda
        csv_map = OrderedDict({"filename": [], "probability": []})
        model.cuda()
        model.eval()
        for i, (input, filepath) in enumerate(tqdm(test_loader)):
            # 3.2 change everything to cuda and get only basename
            filepath = [os.path.basename(x) for x in filepath]
            with torch.no_grad():
                image_var = Variable(input).cuda()
                # 3.3.output
                # print(filepath)
                # print(input,input.shape)
                y_pred = model(image_var)
                print(y_pred.shape)
                smax = nn.Softmax(1)
                smax_out = smax(y_pred)
            # 3.4 save probability to csv files
            csv_map["filename"].extend(filepath)
            for output in smax_out:
                prob = ";".join([str(i) for i in output.data.tolist()])
                csv_map["probability"].append(prob)
        result = pd.DataFrame(csv_map)
        result["probability"] = result["probability"].map(lambda x: [float(i) for i in x.split(";")])
        result.to_csv("./submit/{}_submission.csv".format(config.model_name + "_" + str(folds)), index=False, header=None)
    
    
    # 4. more details to build main function
    def main():
        fold = 0
        # 4.1 mkdirs
        if not os.path.exists(config.submit):
            os.mkdir(config.submit)
        if not os.path.exists(config.weights):
            os.mkdir(config.weights)
        if not os.path.exists(config.best_models):
            os.mkdir(config.best_models)
        if not os.path.exists(config.logs):
            os.mkdir(config.logs)
        if not os.path.exists(config.weights + config.model_name + os.sep + str(fold) + os.sep):
            os.makedirs(config.weights + config.model_name + os.sep + str(fold) + os.sep)
        if not os.path.exists(config.best_models + config.model_name + os.sep + str(fold) + os.sep):
            os.makedirs(config.best_models + config.model_name + os.sep + str(fold) + os.sep)
            # 4.2 get model and optimizer
        model = get_net()
        # model = torch.nn.DataParallel(model)
        model.cuda()
    
        # optimizer = optim.SGD(model.parameters(),lr = config.lr,momentum=0.9,weight_decay=config.weight_decay)
        optimizer = optim.Adam(model.parameters(), lr=config.lr, amsgrad=True, weight_decay=config.weight_decay)
        criterion = nn.CrossEntropyLoss().cuda()
        # 4.3 some parameters for  K-fold and restart model
        start_epoch = 0
        best_precision1 = 0
        best_precision_save = 0
        resume = False
    
        # 4.4 restart the training process
        if resume:
            checkpoint = torch.load(config.best_models + str(fold) + "/model_best.pth.tar")
            start_epoch = checkpoint["epoch"]
            fold = checkpoint["fold"]
            best_precision1 = checkpoint["best_precision1"]
            model.load_state_dict(checkpoint["state_dict"])
            optimizer.load_state_dict(checkpoint["optimizer"])
    
        # 4.5 get files and split for K-fold dataset
        # 4.5.1 read files
        train_data_list = get_files(config.train_data, "train")
        val_data_list = get_files(config.val_data, "val")
        # test_files = get_files(config.test_data,"test")
    
        """ 
        #4.5.2 split
        split_fold = StratifiedKFold(n_splits=3)
        folds_indexes = split_fold.split(X=origin_files["filename"],y=origin_files["label"])
        folds_indexes = np.array(list(folds_indexes))
        fold_index = folds_indexes[fold]
        #4.5.3 using fold index to split for train data and val data
        train_data_list = pd.concat([origin_files["filename"][fold_index[0]],origin_files["label"][fold_index[0]]],axis=1)
        val_data_list = pd.concat([origin_files["filename"][fold_index[1]],origin_files["label"][fold_index[1]]],axis=1)
        """
        # train_data_list,val_data_list = train_test_split(origin_files,test_size = 0.1,stratify=origin_files["label"])
        # 4.5.4 load dataset
        train_dataloader = DataLoader(ChaojieDataset(train_data_list), batch_size=config.batch_size, shuffle=True,
                                      collate_fn=collate_fn, pin_memory=True, num_workers=4)
        val_dataloader = DataLoader(ChaojieDataset(val_data_list, train=False), batch_size=config.batch_size * 2,
                                    shuffle=True, collate_fn=collate_fn, pin_memory=False, num_workers=4)
        # test_dataloader = DataLoader(ChaojieDataset(test_files,test=True),batch_size=1,shuffle=False,pin_memory=False)
        # scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,"max",verbose=1,patience=3)
        scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
        # 4.5.5.1 define metrics
        train_losses = AverageMeter()
        train_top1 = AverageMeter()
        valid_loss = [np.inf, 0, 0]
        model.train()
    
        # 4.5.5 train
        start = timer()
        for epoch in range(start_epoch, config.epochs):
            scheduler.step(epoch)
            train_progressor = ProgressBar(mode="Train", epoch=epoch, total_epoch=config.epochs,
                                           model_name=config.model_name, total=len(train_dataloader))
            # train
            # global iter
            for iter, (input, target) in enumerate(train_dataloader):
                # 4.5.5 switch to continue train process
                train_progressor.current = iter
                model.train()
                input = Variable(input).cuda()
                target = Variable(torch.from_numpy(np.array(target)).long()).cuda()
                # target = Variable(target).cuda()
                output = model(input)
                loss = criterion(output, target)
    
                precision1_train, precision2_train = accuracy(output, target, topk=(1, 2))
                train_losses.update(loss.item(), input.size(0))
                train_top1.update(precision1_train[0], input.size(0))
                train_progressor.current_loss = train_losses.avg
                train_progressor.current_top1 = train_top1.avg
                # backward
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                train_progressor()
            train_progressor.done()
            # evaluate
            lr = get_learning_rate(optimizer)
            # evaluate every half epoch
            valid_loss = evaluate(val_dataloader, model, criterion, epoch)
            is_best = valid_loss[1] > best_precision1
            best_precision1 = max(valid_loss[1], best_precision1)
            try:
                best_precision_save = best_precision1.cpu().data.numpy()
            except:
                pass
            save_checkpoint({
                "epoch": epoch + 1,
                "model_name": config.model_name,
                "state_dict": model.state_dict(),
                "best_precision1": best_precision1,
                "optimizer": optimizer.state_dict(),
                "fold": fold,
                "valid_loss": valid_loss,
            }, is_best, fold)
    
    
    if __name__ == "__main__":
        main()

### 测试函数:test.py ###

# -*- coding: utf-8 -*-
    # @Time    : 2021/1/21 10:06
    # @Author  : Johnson
    import os
    import random
    import time
    import json
    import torch
    import torchvision
    import numpy as np
    import pandas as pd
    import warnings
    from datetime import datetime
    from torch import nn,optim
    from config import config
    from collections import OrderedDict
    from torch.autograd import Variable
    from torch.utils.data import DataLoader
    from dataset.dataloader import *
    from sklearn.model_selection import train_test_split,StratifiedKFold
    from timeit import default_timer as timer
    from models.model import *
    from utils import *
    from IPython import embed
    #1. set random.seed and cudnn performance
    random.seed(config.seed)
    np.random.seed(config.seed)
    torch.manual_seed(config.seed)
    torch.cuda.manual_seed_all(config.seed)
    os.environ["CUDA_VISIBLE_DEVICES"] = config.gpus
    torch.backends.cudnn.benchmark = True
    warnings.filterwarnings('ignore')
    
    #3. test model on public dataset and save the probability matrix
    def test(test_loader,model,folds):
        #3.1 confirm the model converted to cuda
        test_labels = open("%s.csv"%config.model_name,"w")
        model.cuda()
        model.eval()
        tta = False
        for i,(input,filepath) in enumerate(tqdm(test_loader)):
            #3.2 change everything to cuda and get only basename
            filepath = [os.path.basename(x) for x in filepath]
            with torch.no_grad():
                image_var = Variable(input).cuda()
                #3.3.output
                #print(filepath)
                #print(input,input.shape)
                if tta == False:
                    y_pred = model(image_var)
                    smax = nn.Softmax(1)
                    smax_out = smax(y_pred)
                    label = np.argmax(smax_out.cpu().data.numpy())
                    test_labels.write(filepath[0]+","+str(label)+"\n")
    
    #4. more details to build main function
    def main():
        fold = 0
        #4.1 mkdirs
        if not os.path.exists(config.submit):
            os.mkdir(config.submit)
        if not os.path.exists(config.weights):
            os.mkdir(config.weights)
        if not os.path.exists(config.best_models):
            os.mkdir(config.best_models)
        if not os.path.exists(config.logs):
            os.mkdir(config.logs)
        if not os.path.exists(config.weights + config.model_name + os.sep +str(fold) + os.sep):
            os.makedirs(config.weights + config.model_name + os.sep +str(fold) + os.sep)
        if not os.path.exists(config.best_models + config.model_name + os.sep +str(fold) + os.sep):
            os.makedirs(config.best_models + config.model_name + os.sep +str(fold) + os.sep)
        #4.2 get model and optimizer
        model = get_net()
        #model = torch.nn.DataParallel(model)
        model.cuda()
        #4.5 get files and split for K-fold dataset
        test_files = get_files(config.test_data,"test")
        #4.5.4 load dataset
        test_dataloader = DataLoader(ChaojieDataset(test_files,test=True),batch_size=1,shuffle=False,pin_memory=False)
        best_model = torch.load("checkpoints/best_model/%s/0/model_best.pth.tar"%config.model_name)
        model.load_state_dict(best_model["state_dict"])
        test(test_dataloader,model,fold)
    
    if __name__ =="__main__":
        main()

### 参考链接 ###

*  [从实例掌握 pytorch 进行图像分类][pytorch]
 *  [pytorch-image-classification][]
 *  [pytorch-classification][]

[pytorch]: http://spytensor.com/index.php/archives/21/?hipebg=1nia31
[pytorch-image-classification]: https://github.com/spytensor/pytorch-image-classification
[pytorch-classification]: https://github.com/bearpaw/pytorch-classification