train.py

import argparse
import copy
import math
import os
import pickle
import sys
import time

import torch.utils.data
import torchvision
from torch import nn
from torch.nn.utils import clip_grad_norm_
from torchvision.models import mnasnet1_0

import misc
from datasets import cifar, mnist, imagenet
from misc import model_snapshot, AverageMeter, validate, apply_weight_decay, load_pretrained_model, array1d_repr
from model.caffelenet import CaffeLeNet
from model.cifar_resnet import ResNet50
from model.cifar_resnet2 import resnet20
from model.imagenet_alexnet import alexnet
from model.imagenet_mobilenetv1 import mobilenetv1
#from proxyless_nas import proxyless_mobile, proxyless_gpu
from util import quantized_sparsify, sparse_quantize, quantize_with_bits, tensor_round

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Prune-Quant training in pytorch')
    parser.add_argument('--dataset', default='mnist', help='dataset used in the experiment')
    parser.add_argument('--data_dir', default='./ILSVRC_CLS', help='dataset dir in this machine')
    parser.add_argument('--arch', '-a', metavar='ARCH', default='caffelenet')
    parser.add_argument('-j', '--workers', default=2, type=int, metavar='N',
                        help='number of data loading workers')
    parser.add_argument('--epochs', default=120, type=int, metavar='N', help='number of total epochs to run')
    parser.add_argument('--warmupT', default=0, type=float, help='number of total iterations for warmup')
    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: 128)')
    parser.add_argument('--lr', '--learning-rate', default=0.025, type=float, metavar='LR',
                        help='initial learning rate')
    parser.add_argument('--lr_sched', default=None, type=str, help='lr scheduler')
    parser.add_argument('--bit_budget', default=4, type=float, help='bit budget ({1,2...,8})')
    parser.add_argument('--nnz_budget', default=0.5, type=float, help='number of nonzero budget (0.0~1.0)')

    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')
    parser.add_argument('--rho', default=0.05, type=float, help='admm hyperparameter rho')
    parser.add_argument('--gclip', default=-1, type=float, help='gradient clip')

    parser.add_argument('--projint', type=int, default=0,
                        help='how many batches to wait before sparse projection of primal weights')
    parser.add_argument('--dualint', type=int, default=0,
                        help='how many batches to wait before updating duplicate and dual weights')

    parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
                        help='evaluate model on validation set')
    parser.add_argument('--pretrain', default=None, help='file to load pretrained model')
    parser.add_argument('--logdir',
                        help='The directory used to save the trained models',
                        default='log/default', type=str)
    parser.add_argument('--save-every', dest='save_every',
                        help='Saves checkpoints at every specified number of epochs',
                        type=int, default=-1)
    parser.add_argument('--mgpu', action='store_true', help='enable using multiple gpus')
    parser.add_argument('--log_interval', type=int, default=-1,
                        help='how many batches to wait before logging training status')
    parser.add_argument('--seed', type=int, default=1, help='random seed')
    parser.add_argument('--eval_tr', action='store_true', help='evaluate training set')
    parser.add_argument('--prox', action='store_true', help='use proximal op for primal update')
    parser.add_argument('--dp', default=0.0, type=float, help='dropout rate')
    parser.add_argument('--quant', action='store_true', help='only perform quantization')
    parser.add_argument('--prune', action='store_true', help='only perform pruning')
    parser.add_argument('--bwlb', type=int, default=1, help='the lower bound of bitwidth')
    parser.add_argument('--bits_epoch', type=int, default=-1, help='maximum epochs allowing update bits')
    parser.add_argument('--kdtemp', default=0.0, type=float, help='knowledge distillation temperature')
    parser.add_argument('--optim', default='sgd', help='optimizer to use')
    parser.add_argument('--fixedbits', action='store_true', help='use fixed bitwidth')

    best_acc = 0
    old_file = None

    args = parser.parse_args()

    args.cuda = torch.cuda.is_available()

    torch.backends.cudnn.benchmark = True

    # set up random seeds
    misc.seed_torch(args.seed)

    # create log file
    args.logdir = os.path.join(os.path.dirname(__file__), args.logdir)

    if os.path.exists(args.logdir):
        ans = misc.query_yes_no('Are you sure to overwrite the original directory: {}?'.format(args.logdir))
        if ans:
            # rm old contents in dir
            print('remove old contents in {}'.format(args.logdir))
            os.system('rm -rf ' + args.logdir)
        else:
            exit()
    misc.logger.init(args.logdir, 'train_log')
    print = misc.logger.info

    print('command:\npython {}'.format(' '.join(sys.argv)))

    print("=================FLAGS==================")
    for k, v in args.__dict__.items():
        print('{}: {}'.format(k, v))
    print("========================================")

    # create model
    teacher_model = None
    if args.dataset == 'cifar10':
        if args.arch == 'resnet50':
            model = ResNet50()
        elif args.arch == 'resnet20':
            model = resnet20()
        else:
            raise NotImplementedError
    elif args.dataset == 'mnist':
        if args.arch == 'caffelenet':
            model = CaffeLeNet()
        else:
            raise NotImplementedError
    elif args.dataset == 'imagenet':
        if args.arch == 'resnet18':
            model = torchvision.models.resnet18(pretrained=args.pretrain == 'pytorch')
        elif args.arch == 'alexnet':
            model = alexnet(pretrained=args.pretrain == 'pytorch', dropout=args.dp)
            if args.kdtemp > 0.0:
                # + knowledge distillation loss
                teacher_model = alexnet(pretrained=True)
                for param in teacher_model.parameters():
                    param.requires_grad = False
                teacher_model.eval()
        elif args.arch == 'mobilenetv1':
            model = mobilenetv1(pretrained=args.pretrain == 'pytorch')
        elif args.arch == 'mnasnet1_0':
            model = mnasnet1_0(pretrained=args.pretrain == 'pytorch')
        elif args.arch == 'proxyless_mobile':
            model = proxyless_mobile(pretrained=args.pretrain == 'pytorch')
        elif args.arch == 'proxyless_gpu':
            model = proxyless_gpu(pretrained=args.pretrain == 'pytorch')
        else:
            raise NotImplementedError
    else:
        raise NotImplementedError

    if not hasattr(model, 'weight_bits'):
        model.register_buffer('weight_bits', torch.tensor([0] * len([m for m in model.modules()
                                                                     if isinstance(m, nn.Conv2d)
                                                                     or isinstance(m, nn.Linear)])))
    else:
        raise ValueError

    # pretrained model
    if args.pretrain != 'pytorch':
        load_pretrained_model(args.pretrain, model, strict=False)

    net_model = model

    if args.mgpu:
        assert len(os.environ['CUDA_VISIBLE_DEVICES'].split(',')) > 1
        if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
            model.features = torch.nn.DataParallel(model.features)
            if teacher_model is not None:
                teacher_model.features = torch.nn.DataParallel(teacher_model.features)
        else:
            model = torch.nn.DataParallel(model)
            if teacher_model is not None:
                teacher_model = torch.nn.DataParallel(teacher_model)

    if args.cuda:
        model.cuda()
        if teacher_model is not None:
            teacher_model.cuda()

    if args.dataset == 'cifar10':
        train_loader, val_loader = cifar.get10(batch_size=args.batch_size, data_root='./.data', train=True, val=True,
                                               num_workers=args.workers)
        train_loader4eval = train_loader
    elif args.dataset == 'mnist':
        train_loader, val_loader = mnist.get(batch_size=args.batch_size, data_root='./.data', train=True, val=True,
                                             num_workers=args.workers)
        train_loader4eval = train_loader
    elif args.dataset == 'imagenet':
        train_loader, val_loader, train_loader4eval = imagenet.get_data_loaders(args.data_dir,
                                                                                batch_size=args.batch_size,
                                                                                val_batch_size=args.batch_size,
                                                                                num_workers=args.workers,
                                                                                nsubset=-1,
                                                                                normalize=None)
    else:
        raise NotImplementedError

    loss_func = lambda m, x, y: misc.cross_entropy(m(x), y)

    if args.optim == 'sgd':
        optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=0.0,
                                    nesterov=False)
    elif args.optim == 'adam':
        optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=0.0)

    warmupT = int(args.warmupT * len(train_loader))

    # lr scheduler setup
    if args.lr_sched is not None:
        train_loader_len = len(train_loader)
        lr_sched = args.lr_sched.split(',')
        if lr_sched[0] == 'cos':
            if len(lr_sched) > 1:
                min_lr = float(lr_sched[1])
            else:
                min_lr = 0.0
            lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
                                                                      train_loader_len * (args.epochs - args.warmupT),
                                                                      eta_min=min_lr,
                                                                      last_epoch=len(
                                                                          train_loader) * args.start_epoch - 1)
        elif lr_sched[0] == 'plat':
            lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=float(lr_sched[1]),
                                                                      threshold=float(lr_sched[2]),
                                                                      patience=2)
        elif lr_sched[0] == 'step':
            lr_milestones = [int(i) for i in lr_sched[1:]]
            print('lr multi-step decay, milestones={}'.format(lr_milestones))
            lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
                                                                gamma=0.5,
                                                                milestones=lr_milestones,
                                                                last_epoch=args.start_epoch - 1)
        elif lr_sched[0] == 'exp':
            factor = float(lr_sched[1])
            print('lr exp decay, factor={}'.format(factor))
            lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=factor)
        else:
            raise NotImplementedError
    else:
        lr_scheduler = None

    model_weights = [p for name, p in model.named_parameters() if name.endswith('weight')]
    conv2d_weights = [m.weight for m in model.modules() if isinstance(m, nn.Conv2d) or isinstance(m, nn.BatchNorm2d)]
    linear_weights = [m.weight for m in model.modules() if isinstance(m, nn.Linear)]

    if args.projint <= 0:
        args.projint = len(train_loader)
    if args.dualint <= 0:
        args.dualint = len(train_loader)
    assert args.dualint == len(
        train_loader) or args.dualint % args.projint == 0, 'dualint should be dividable by projint'
    conv_weights = [m.weight for m in model.modules() if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear)]
    n_conv_layers = len(conv_weights)
    if not args.prune:
        weight_bits = [math.floor(args.bit_budget)] * n_conv_layers
    else:
        assert not args.quant
        args.bit_budget = 32
        weight_bits = [32] * n_conv_layers
    num_weights = [w.numel() for w in conv_weights]
    print('number of weights: {}'.format(num_weights))
    if not args.quant:
        num_nnz = [math.ceil(args.nnz_budget * w.numel()) for w in conv_weights]
    else:
        num_nnz = [num for num in num_weights]

    model_size = sum([args.bit_budget * num_nnz[i] for i in range(n_conv_layers)])
    model_size_lb = sum([weight_bits[i] * num_nnz[i] for i in range(n_conv_layers)])

    full_model_size = sum([w.numel() * 32 for w in conv_weights])
    print('target model size={:.4e} bits / {:.4e} bits (compression rate:{:.4e})'.format(model_size,
                                                                                         full_model_size,
                                                                                         float(
                                                                                             model_size) / full_model_size))

    assert args.bwlb <= min(weight_bits)

    for i in range(len(weight_bits)):
        net_model.weight_bits[i] = weight_bits[i]

    if args.evaluate:
        validate(val_loader, model, loss_func=loss_func)
        exit()

    # tensors used in ADMM algorithm
    if not args.prune:
        conv_weights_dup = [w.data.clone() for w in conv_weights]
        conv_weights_dual = [torch.zeros_like(w.data) for w in conv_weights]
    residual = [0.0] * n_conv_layers
    residual2 = [0.0] * n_conv_layers

    reserved_cluster = 0 if args.quant else 1
    log_tic = time.time()
    losses = AverageMeter()
    model_eval = copy.deepcopy(model)
    model_eval_conv_weights = [m.weight for m in model_eval.modules()
                               if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear)]

    loss_func = lambda m, x, y: misc.classify_loss(m, x, y, teacher_model, args.kdtemp)

    if args.log_interval <= 0:
        args.log_interval = len(train_loader)
    tr_loss_hist = []
    tr_loss_hist_raw = []
    ts_loss_hist = []
    residual_hist = []
    for epoch in range(args.start_epoch, args.epochs):
        losses.reset()
        gclip_time = 0
        # train for one epoch
        print('current lr {:.5e}'.format(optimizer.param_groups[0]['lr']))
        for batch_idx, (data, target) in enumerate(train_loader):
            # lr schedule
            t = float(batch_idx + epoch * len(train_loader))
            if t < warmupT:
                lr = min(1.0, (t + 1) / float(warmupT)) * args.lr
                for param_group in optimizer.param_groups:
                    param_group['lr'] = lr

            model.train()
            if args.cuda:
                data, target = data.cuda(), target.cuda()

            w_loss = loss_func(model, data, target)

            if args.prune:
                l2_aug_loss = 0.0
            else:
                l2_aug_loss = 0.5 * args.rho * \
                              sum([torch.sum((conv_weights[i] - conv_weights_dup[i] + conv_weights_dual[i]) ** 2)
                                   for i in range(n_conv_layers)])
                if args.prox:
                    l2_aug_loss = l2_aug_loss.data.item()
            primal_loss = w_loss + l2_aug_loss
            # losses stats
            losses.update(primal_loss.item(), data.size(0))
            # update network weights
            optimizer.zero_grad()
            primal_loss.backward()
            # apply weight_decay
            apply_weight_decay(model_weights, args.weight_decay)
            # gradient norm clip
            if args.gclip > 0:
                total_norm = clip_grad_norm_(model.parameters(), args.gclip, norm_type=float('inf'))
                if total_norm > args.gclip:
                    gclip_time += 1
            optimizer.step()

            if (not args.prune) and args.prox:
                lr = optimizer.param_groups[0]['lr']
                beta = lr * args.rho
                for i, p in enumerate(conv_weights):
                    p.data.add_(beta, conv_weights_dup[i] - conv_weights_dual[i])
                    p.data /= (1.0 + beta)

            if (args.pretrain is not None and batch_idx == 0 and epoch == 0) or \
                    ((batch_idx + 1) % args.projint == 0 or batch_idx == len(train_loader) - 1):
                # sparse projection (in-place)
                if not args.quant:
                    num_nnz = quantized_sparsify(conv_weights, weight_bits, model_size_lb, in_place=True)[1]
                    model_size_lb = model_size
                else:
                    num_nnz = [num for num in num_weights]

            if not args.prune:
                if (args.pretrain is not None and batch_idx == 0 and epoch == 0) or \
                        ((batch_idx + 1) % args.dualint == 0 or batch_idx == len(train_loader) - 1):
                    # Update conv_weights_dup (quantization projection)
                    if args.fixedbits or (args.bits_epoch > 0 and epoch > args.bits_epoch):
                        res_weight_list = \
                            quantize_with_bits([conv_weights[i].data + conv_weights_dual[i] for i in range(n_conv_layers)],
                                               weight_bits, in_place=False, dictnz=args.quant)[0]
                    else:
                        res_weight_list, weight_bits, offered_cluster = \
                            sparse_quantize([conv_weights[i].data + conv_weights_dual[i] for i in range(n_conv_layers)],
                                            num_nnz, model_size, include_dict=False, in_place=False, dictnz=args.quant,
                                            bwlb=args.bwlb)

                    for i in range(len(weight_bits)):
                        net_model.weight_bits[i] = weight_bits[i]

                    for i in range(n_conv_layers):
                        conv_weights_dup[i].copy_(res_weight_list[i])

                    # Update dual weights
                    for i in range(n_conv_layers):
                        diff = conv_weights[i].data - conv_weights_dup[i]
                        conv_weights_dual[i] += diff
                        residual[i] = (diff ** 2).sum().item()

            if epoch >= args.warmupT and lr_scheduler is not None:
                # increment lr_scheduler each epoch or cosine decay each iteration
                if batch_idx == len(train_loader) - 1 or lr_sched[0] == 'cos':
                    lr_scheduler.step()

            if (batch_idx + 1) % args.log_interval == 0 or batch_idx == len(train_loader) - 1:
                print('+-------------- epoch {}, batch {}/{} ----------------+'.format(epoch, batch_idx + 1,
                                                                                       len(train_loader)))
                log_toc = time.time()
                print('Primal update: Loss={:.4f} (losses_avg={:.4f})'
                      ', lr={:.4e}, time_elapsed={:.3f}s'.format(
                    losses.val, losses.avg, optimizer.param_groups[0]['lr'], log_toc - log_tic))
                if args.gclip > 0:
                    print('gclip times={}'.format(gclip_time))
                # print(layers_stat(model, param_names='weight', param_filter=lambda p: p.dim() > 1))
                print('Dual update: rho={:.4f}'.format(args.rho))
                print('residual=\t{}'.format(array1d_repr(residual)))
                print('residual2=\t{}'.format(array1d_repr(residual2)))

                print('num_nnz=\t{}'.format(array1d_repr([num_nnz[i] / conv_weights[i].numel()
                                                          for i in range(n_conv_layers)], format='{:.3f}')))
                print('weight_bits=\t{}'.format(array1d_repr(weight_bits, format='{:.0f}')))

                log_tic = time.time()
                print('+-----------------------------------------------------+')
        # evaluate
        model_eval.load_state_dict(model.state_dict())
        if not args.prune:
            quantize_with_bits(model_eval_conv_weights, weight_bits, in_place=True, dictnz=args.quant)

        num_nnz_eval = [float(m.weight.nonzero().shape[0]) for m in model_eval.modules()
                        if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear)]
        weight_bits_eval = [math.log2(max(1.0, tensor_round(m.weight.data, n=6).unique().shape[0] - reserved_cluster))
                            for m in model_eval.modules() if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear)]
        print('n_nnz_eval=\t{}'.format(num_nnz_eval))
        print('num_nnz_eval=\t{}'.format(array1d_repr([num_nnz_eval[i] / conv_weights[i].numel()
                                                       for i in range(n_conv_layers)], format='{:.3f}')))
        print('nnz_eval=\t{:.4e}'.format(sum(num_nnz_eval) / sum([conv_weights[i].numel()
                                                                  for i in range(n_conv_layers)])))
        print('weight_bits_eval=\t{}'.format(array1d_repr(weight_bits_eval, format='{:.0f}')))
        print('ave_weight_bits_eval=\t{:.4e}'.format(sum([num_nnz_eval[i] * weight_bits_eval[i]
                                                          for i in range(len(num_nnz_eval))]) / sum(num_nnz_eval)))
        print(misc.layers_stat(model, param_names='weight', param_filter=lambda p: p.dim() > 1))
        if args.eval_tr:
            print('training set:')
            tr_loss_hist.append(validate(train_loader4eval, model_eval, loss_func=loss_func)[1])
            tr_loss_hist_raw.append(validate(train_loader4eval, model, loss_func=loss_func, verbose=False)[1])
        residual_hist.append(sum(residual) / float(sum(num_weights)))
        print('test set:')
        prec1, ts_loss = validate(val_loader, model_eval, loss_func=loss_func)
        ts_loss_hist.append(ts_loss)
        print('current model size={:.4e} bits'.format(
            sum([weight_bits_eval[i] * num_nnz_eval[i] for i in range(n_conv_layers)])))
        print(
            'compression rate={:.4e}'.format(sum([weight_bits_eval[i] * num_nnz_eval[i] for i in range(n_conv_layers)])
                                             / full_model_size))
        print('compression rate={:.4e} (train)'
              .format(sum([weight_bits[i] * num_nnz[i] for i in range(n_conv_layers)]) / full_model_size))
        print('======================================================')

        # remember best prec@1 and save checkpoint
        if prec1 > best_acc:
            print('find accuracy {:4f} > {:.4f}'.format(prec1, best_acc))
            new_file = os.path.join(args.logdir, 'model_best-{}.pkl'.format(epoch))
            misc.model_snapshot(model_eval, new_file, old_file=old_file, verbose=True)
            best_acc = prec1
            old_file = new_file

        if epoch > 0 and args.save_every > 0 and epoch % args.save_every == 0:
            model_snapshot(model, os.path.join(args.logdir, 'model_epoch{}.pt'.format(epoch)))

        # save the lastest model
        model_snapshot(model, os.path.join(args.logdir, 'model_latest.pt'))

    # print(tr_loss_hist)
    # print(ts_loss_hist)
    # print(residual_hist)
    with open(os.path.join(args.logdir, 'loss_hist.pkl'), 'wb') as f:
        pickle.dump({'tr_loss_raw': tr_loss_hist_raw, 'tr_loss': tr_loss_hist,
                     'ts_loss': ts_loss_hist, 'residual': residual_hist}, f)