PyTorch代码学习-ImageNET训练

文章说明：本人学习pytorch/examples/ImageNET/main()理解（待续）

# -*- coding: utf-8 -*-import argparse  # 命令行解释器相关程序，命令行解释器import os        # 操作系统文件相关import shutil    # 文件高级操作import time      # 调用时间模块import torchimport torch.nn as nnimport torch.nn.parallelimport torch.backends.cudnn as cudnn        # gpu 使用import torch.distributed as dist            # 分布式（pytorch 0.2)import torch.optim                          # 优化器import torch.utils.dataimport torch.utils.data.distributedimport torchvision.transforms as transformsimport torchvision.datasets as datasetsimport torchvision.models as models# name中若为小写且不以‘——’开头，则对其进行升序排列model_names = sorted(name for name in models.__dict__    if name.islower() and not name.startswith("__")    and callable(models.__dict__[name]))                    # callable功能为判断返回对象是否可调用（即某种功能）。# 创建argparse.ArgumentParser对象parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')# 添加命令行元素parser.add_argument('data', metavar='DIR',                    help='path to dataset')parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18',                    choices=model_names,                    help='model architecture: ' +                        ' | '.join(model_names) +                        ' (default: resnet18)')parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',                    help='number of data loading workers (default: 4)')parser.add_argument('--epochs', default=90, type=int, metavar='N',                    help='number of total epochs to run')parser.add_argument('--start-epoch', default=0, type=int, metavar='N',                    help='manual epoch number (useful on restarts)')parser.add_argument('-b', '--batch-size', default=256, type=int,                    metavar='N', help='mini-batch size (default: 256)')parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,                    metavar='LR', help='initial learning rate')parser.add_argument('--momentum', default=0.9, type=float, metavar='M',                    help='momentum')parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,                    metavar='W', help='weight decay (default: 1e-4)')parser.add_argument('--print-freq', '-p', default=10, type=int,                    metavar='N', help='print frequency (default: 10)')parser.add_argument('--resume', default='', type=str, metavar='PATH',                    help='path to latest checkpoint (default: none)')parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',                    help='evaluate model on validation set')parser.add_argument('--pretrained', dest='pretrained', action='store_true',                    help='use pre-trained model')parser.add_argument('--world-size', default=1, type=int,                    help='number of distributed processes')parser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str,                    help='url used to set up distributed training')parser.add_argument('--dist-backend', default='gloo', type=str,                    help='distributed backend')# 定义参数best_prec1 = 0# 定义主函数main()def main():    global args, best_prec1    # 使用函数parse_args()进行参数解析，输入默认是sys.argv[1:]，    # 返回值是一个包含命令参数的Namespace，所有参数以属性的形式存在，比如args.myoption。    args = parser.parse_args()########## 使用多播地址进行初始化    args.distributed = args.world_size > 1    if args.distributed:        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,                                world_size=args.world_size)##### step1: create model and set GPU     # 导入pretrained model 或者创建model    if args.pretrained:        # format 格式化表达字符串，上述默认arch为resnet18        print("=> using pre-trained model '{}'".format(args.arch))              model = models.__dict__[args.arch](pretrained=True)    else:        print("=> creating model '{}'".format(args.arch))        model = models.__dict__[args.arch]()    # 分布式运行，可实现在多块GPU上运行    if not args.distributed:        if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):            # 批处理，多GPU默认用dataparallel使用在多块gpu上            model.features = torch.nn.DataParallel(model.features)                       model.cuda()        else:            model = torch.nn.DataParallel(model).cuda()    else:        # Wrap model in DistributedDataParallel (CUDA only for the moment)        model.cuda()        model = torch.nn.parallel.DistributedDataParallel(model)##### step2: define loss function (criterion) and optimizer    # 使用交叉熵损失函数    criterion = nn.CrossEntropyLoss().cuda()                                # optimizer 使用 SGD + momentum    # 动量，默认设置为0.9    optimizer = torch.optim.SGD(model.parameters(), args.lr,                                momentum=args.momentum,                                # 权值衰减，默认为1e-4                                                 weight_decay=args.weight_decay)            # 恢复模型（详见模型存取与恢复）####step3：optionally resume from a checkpoint    if args.resume:        if os.path.isfile(args.resume):                                 # 判断返回的是不是文件            print("=> loading checkpoint '{}'".format(args.resume))            checkpoint = torch.load(args.resume)                        # load 一个save的对象            args.start_epoch = checkpoint['epoch']                      # default = 90            best_prec1 = checkpoint['best_prec1']                       # best_prec1 = 0            model.load_state_dict(checkpoint['state_dict'])            optimizer.load_state_dict(checkpoint['optimizer'])          # load_state_dict:恢复模型            print("=> loaded checkpoint '{}' (epoch {})"                  .format(args.resume, checkpoint['epoch']))        else:            print("=> no checkpoint found at '{}'".format(args.resume))    cudnn.benchmark = True##### step4: Data loading code base of dataset(have downloaded) and normalize    # 从 train、val文件中导入数据    traindir = os.path.join(args.data, 'train')    valdir = os.path.join(args.data, 'val')    # 数据预处理：normalize: - mean / std    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],                                            std=[0.229, 0.224, 0.225])    # ImageFolder 一个通用的数据加载器    train_dataset = datasets.ImageFolder(        traindir,        # 对数据进行预处理        transforms.Compose([                      # 将几个transforms 组合在一起            transforms.RandomSizedCrop(224),      # 随机切再resize成给定的size大小            transforms.RandomHorizontalFlip(),    # 概率为0.5，随机水平翻转。            transforms.ToTensor(),                # 把一个取值范围是[0,255]或者shape为(H,W,C)的numpy.ndarray，                                                  # 转换成形状为[C,H,W]，取值范围是[0,1.0]的torch.FloadTensor            normalize,        ]))#######    if args.distributed:        # Use a DistributedSampler to restrict each process to a distinct subset of the dataset.        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)    else:        train_sampler = None######    # train 数据下载及预处理    train_loader = torch.utils.data.DataLoader(        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),        num_workers=args.workers, pin_memory=True, sampler=train_sampler)    val_loader = torch.utils.data.DataLoader(        datasets.ImageFolder(valdir, transforms.Compose([             # 重新改变大小为`size`，若：height>width`,则：(size*height/width, size)            transforms.Scale(256),            # 将给定的数据进行中心切割，得到给定的size。            transforms.CenterCrop(224),            transforms.ToTensor(),            normalize,        ])),        batch_size=args.batch_size, shuffle=False,        num_workers=args.workers, pin_memory=True)         # default workers = 4##### step5: 验证函数    if args.evaluate:        validate(val_loader, model, criterion)             # 自定义的validate函数，见下        return##### step6:开始训练模型    for epoch in range(args.start_epoch, args.epochs):        # Use .set_epoch() method to reshuffle the dataset partition at every iteration        if args.distributed:            train_sampler.set_epoch(epoch)        adjust_learning_rate(optimizer, epoch)      # adjust_learning_rate 自定义的函数，见下        # train for one epoch        train(train_loader, model, criterion, optimizer, epoch)        # evaluate on validation set        prec1 = validate(val_loader, model, criterion)        # remember best prec@1 and save checkpoint        is_best = prec1 > best_prec1        best_prec1 = max(prec1, best_prec1)        save_checkpoint({            'epoch': epoch + 1,            'arch': args.arch,            'state_dict': model.state_dict(),            'best_prec1': best_prec1,            'optimizer' : optimizer.state_dict(),        }, is_best)# 定义相关函数# def train 函数def train(train_loader, model, criterion, optimizer, epoch):    batch_time = AverageMeter()    data_time = AverageMeter()    losses = AverageMeter()    top1 = AverageMeter()    top5 = AverageMeter()    # switch to train mode    model.train()    end = time.time()    for i, (input, target) in enumerate(train_loader):        # measure data loading time        data_time.update(time.time() - end)        target = target.cuda(async=True)        input_var = torch.autograd.Variable(input)        target_var = torch.autograd.Variable(target)        # compute output        output = model(input_var)        # criterion 为定义过的损失函数        loss = criterion(output, target_var)                # measure accuracy and record loss        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))        losses.update(loss.data[0], input.size(0))        top1.update(prec1[0], input.size(0))        top5.update(prec5[0], input.size(0))        # compute gradient and do SGD step        optimizer.zero_grad()        loss.backward()        optimizer.step()        # measure elapsed time        batch_time.update(time.time() - end)        end = time.time()        # 每十步输出一次        if i % args.print_freq == 0:     # default=10            print('Epoch: [{0}][{1}/{2}]\t'                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(                   epoch, i, len(train_loader), batch_time=batch_time,                   data_time=data_time, loss=losses, top1=top1, top5=top5))def validate(val_loader, model, criterion):    batch_time = AverageMeter()    losses = AverageMeter()    top1 = AverageMeter()    top5 = AverageMeter()    # switch to evaluate mode    model.eval()    end = time.time()    for i, (input, target) in enumerate(val_loader):        target = target.cuda(async=True)        # 这是一种用来包裹张量并记录应用的操作        """        Attributes:        data: 任意类型的封装好的张量。        grad: 保存与data类型和位置相匹配的梯度，此属性难以分配并且不能重新分配。        requires_grad: 标记变量是否已经由一个需要调用到此变量的子图创建的bool值。只能在叶子变量上进行修改。        volatile: 标记变量是否能在推理模式下应用（如不保存历史记录）的bool值。只能在叶变量上更改。        is_leaf: 标记变量是否是图叶子(如由用户创建的变量)的bool值.        grad_fn: Gradient function graph trace.        Parameters:        data (any tensor class): 要包装的张量.        requires_grad (bool): bool型的标记值. **Keyword only.**        volatile (bool): bool型的标记值. **Keyword only.**        """        input_var = torch.autograd.Variable(input, volatile=True)        target_var = torch.autograd.Variable(target, volatile=True)        # compute output        output = model(input_var)        loss = criterion(output, target_var)        # measure accuracy and record loss        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))        losses.update(loss.data[0], input.size(0))        top1.update(prec1[0], input.size(0))        top5.update(prec5[0], input.size(0))        # measure elapsed time        batch_time.update(time.time() - end)        end = time.time()        if i % args.print_freq == 0:            print('Test: [{0}/{1}]\t'                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(                   i, len(val_loader), batch_time=batch_time, loss=losses,                   top1=top1, top5=top5))    print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'          .format(top1=top1, top5=top5))    return top1.avg# 保存当前节点def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):    torch.save(state, filename)    if is_best:        shutil.copyfile(filename, 'model_best.pth.tar')# 计算并存储参数当前值或平均值class AverageMeter(object):    # Computes and stores the average and current value    """       batch_time = AverageMeter()       即 self = batch_time       则 batch_time 具有__init__，reset，update三个属性，       直接使用batch_time.update()调用       功能为：batch_time.update(time.time() - end)               仅一个参数，则直接保存参数值        对应定义：def update(self, val, n=1)        losses.update(loss.data[0], input.size(0))        top1.update(prec1[0], input.size(0))        top5.update(prec5[0], input.size(0))        这些有两个参数则求参数val的均值，保存在avg中##不确定##    """    def __init__(self):        self.reset()       # __init__():reset parameters    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# 更新 learning_rate ：每30步，学习率降至前的10分之1def adjust_learning_rate(optimizer, epoch):    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""    lr = args.lr * (0.1 ** (epoch // 30))            # args.lr = 0.1 , 即每30步，lr = lr /10    for param_group in optimizer.param_groups:       # 将更新的lr 送入优化器 optimizer 中，进行下一次优化        param_group['lr'] = lr# 计算准确度def accuracy(output, target, topk=(1,)):    """Computes the precision@k for the specified values of k    prec1, prec5 = accuracy(output.data, target, topk=(1, 5))    """    maxk = max(topk)    # size函数：总元素的个数    batch_size = target.size(0)    # topk函数选取output前k大个数    _, pred = output.topk(maxk, 1, True, True)    ##########不了解t()    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 resif __name__ == '__main__':    main()