quant_train.py

import wandb

import argparse
import os
import random
import shutil
import time
import logging
import warnings

import torch
torch.manual_seed(0) #  FIX random sampler on training data 

import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import torch.optim
import torch.multiprocessing as mp
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets

from src.data_utils import AverageMeter, ProgressMeter
from src.post_quant.cle import cle_for_resmlp, cle_for_resmlp_v3, cle_for_resmlp_v4
from src.post_quant.rca import rca_for_resmlp
from src.models import *
from timm.scheduler.cosine_lr import CosineLRScheduler

parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
parser.add_argument('--data', metavar='DIR',
                    help='path to dataset')
parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet18',
                    help='model architecture')
parser.add_argument('--teacher-arch',
                    type=str,
                    default='resnet101',
                    help='teacher network used to do distillation')
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), this is the total '
                         'batch size of all GPUs on the current node when '
                         'using Data Parallel or Distributed Data Parallel')
parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
                    metavar='LR', help='initial learning rate', dest='lr')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum')
parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)',
                    dest='weight_decay')
parser.add_argument('-p', '--print-freq', 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 nodes for distributed training')
parser.add_argument('--rank', default=-1, type=int,
                    help='node rank for distributed training')
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='nccl', type=str,
                    help='distributed backend')
parser.add_argument('--seed', default=None, type=int,
                    help='seed for initializing training. ')
parser.add_argument('--gpu', default=None, type=int,
                    help='GPU id to use.')
parser.add_argument('--multiprocessing-distributed', action='store_true',
                    help='Use multi-processing distributed training to launch '
                         'N processes per node, which has N GPUs. This is the '
                         'fastest way to use PyTorch for either single node or '
                         'multi node data parallel training')
parser.add_argument('--act-range-momentum',
                    type=float,
                    default=-1,
                    help='momentum of the activation range moving average, '
                         '-1 stands for using minimum of min and maximum of max')
parser.add_argument('--quant-mode',
                    type=str,
                    default='symmetric',
                    choices=['asymmetric', 'symmetric'],
                    help='quantization mode')
parser.add_argument('--save-path',
                    type=str,
                    default='checkpoints/imagenet/test/',
                    help='path to save the quantized model')
parser.add_argument('--data-percentage',
                    type=float,
                    default=1,
                    help='data percentage of training data')
parser.add_argument('--fix-BN',
                    action='store_true',
                    help='whether to fix BN statistics and fold BN during training')
parser.add_argument('--fix-BN-threshold',
                    type=int,
                    default=None,
                    help='when to start training with fixed and folded BN,'
                         'after the threshold iteration, the original fix-BN will be overwritten to be True')
parser.add_argument('--evaluate-times',
                    type=int,
                    default=-1,
                    help='The number of evaluations during one epoch')
parser.add_argument('--quant-scheme',
                    type=str,
                    default='uniform8',
                    help='quantization bit configuration')
parser.add_argument('--resume-quantize',
                    action='store_true',
                    help='if True map the checkpoint to a quantized model,'
                         'otherwise map the checkpoint to an ordinary model and then quantize')
parser.add_argument('--act-percentile',
                    type=float,
                    default=0,
                    help='the percentage used for activation percentile'
                         '(0 means no percentile, 99.9 means cut off 0.1%)')
parser.add_argument('--weight-percentile',
                    type=float,
                    default=0,
                    help='the percentage used for weight percentile'
                         '(0 means no percentile, 99.9 means cut off 0.1%)')
parser.add_argument('--bias-bit',
                    type=int,
                    default=32,
                    help='quantizaiton bit-width for bias')
parser.add_argument('--distill-method',
                    type=str,
                    default='None',
                    help='you can choose None or KD_naive')
parser.add_argument('--distill-alpha',
                    type=float,
                    default=0.95,
                    help='how large is the ratio of normal loss and teacher loss')
parser.add_argument('--temperature',
                    type=float,
                    default=6,
                    help='how large is the temperature factor for distillation')
parser.add_argument('--skip-connection-fp',
                    action='store_true',
                    help='whether to ignore quantizating skip-connecitions')
parser.add_argument('--no-quant',
                    action='store_true',
                    help='if set to true, run model with quantization all disabled')
parser.add_argument('--regular',
                    action='store_true',
                    help='if set to true, run with original model')
parser.add_argument('--cle',
                    action='store_true',
                    help='if set to true, run cle before QAT')  
parser.add_argument('--freeze-w',
                    action='store_true',
                    help='if set to true, freeze weight and update scale only') 
parser.add_argument('--skip-val',
                    action='store_true',
                    help='if set to true, skip validation')       
parser.add_argument('--rca',
                    action='store_true',
                    help='if set to true, run rca before QAT')   
parser.add_argument('--wandb',
                    action='store_true',
                    help='if set to true, log with wandb')
parser.add_argument('--inner-wd', default=1e-4, type=float,
                    metavar='IN_W', help='weight decay for inner bias (default: 1e-4)')
parser.add_argument('--load_pretrain', default='', type=str, metavar='PATH',
                    help='path to latest fp32 checkpoint (default: none)')
best_acc1 = 0

arch_dict = {'q_resmlp': resmlp_24, 'q_resmlp_norm': resmlp_24_norm, 'q_resmlp_v2': resmlp_24, 'q_resmlp_v3': resmlp_24, 'q_resmlp_v4': resmlp_24_v4}
quantize_arch_dict = {'q_resmlp': q_resmlp, 'q_resmlp_norm': q_resmlp_norm, 'q_resmlp_v2': q_resmlp_v2, 'q_resmlp_v3': q_resmlp_v3, 'q_resmlp_v4': q_resmlp_v4}

args = parser.parse_args()
if not os.path.exists(args.save_path):
    os.makedirs(args.save_path)

logging.basicConfig(format='%(asctime)s - %(message)s',
                    datefmt='%d-%b-%y %H:%M:%S', filename=args.save_path + 'log.log')
logging.getLogger().setLevel(logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler())

logging.info(args)

print(args.batch_size)
def main():
    if args.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        cudnn.deterministic = True
        warnings.warn('You have chosen to seed training. '
                      'This will turn on the CUDNN deterministic setting, '
                      'which can slow down your training considerably! '
                      'You may see unexpected behavior when restarting '
                      'from checkpoints.')

    if args.gpu is not None:
        warnings.warn('You have chosen a specific GPU. This will completely '
                      'disable data parallelism.')

    if args.dist_url == "env://" and args.world_size == -1:
        args.world_size = int(os.environ["WORLD_SIZE"])

    args.distributed = args.world_size > 1 or args.multiprocessing_distributed

    ngpus_per_node = torch.cuda.device_count()
    if args.multiprocessing_distributed:
        # Since we have ngpus_per_node processes per node, the total world_size
        # needs to be adjusted accordingly
        args.world_size = ngpus_per_node * args.world_size
        # Use torch.multiprocessing.spawn to launch distributed processes: the
        # main_worker process function
        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
    else:
        # Simply call main_worker function
        main_worker(args.gpu, ngpus_per_node, args)


def main_worker(gpu, ngpus_per_node, args):
    global best_acc1
    args.gpu = gpu

    if args.gpu is not None:
        logging.info("Use GPU: {} for training".format(args.gpu))

    # create model
    if args.pretrained and not args.resume:
        logging.info("=> using pre-trained model '{}'".format(args.arch))
        arch = arch_dict[args.arch]
        model = arch(pretrained=True)

    else:
        logging.info("=> creating model '{}'".format(args.arch))
        arch = arch_dict[args.arch]
        model = arch(pretrained=False)
    
    if args.load_pretrain:
        logging.info("=> loading fp32 checkpoint '{}'".format(args.load_pretrain))
        model.load_state_dict(torch.load(args.load_pretrain)['model'])

    if args.cle:
        logging.info("=> Applying CLE on model")
        if args.arch == "q_resmlp_v4":
            cle_for_resmlp_v4(model)
        else:
            cle_for_resmlp(model)
    

    # TODO: rca calibration
    # if args.rca:
    #     logging.info("=> Applying RCA on model")
    #     rca_for_resmlp(model)
        

    # if args.resume and not args.resume_quantize:
    #     if os.path.isfile(args.resume):
    #         logging.info("=> loading checkpoint '{}'".format(args.resume))

    #         checkpoint = torch.load(args.resume)['state_dict']
    #         model_key_list = list(model.state_dict().keys())
    #         for key in model_key_list:
    #             if 'num_batches_tracked' in key: model_key_list.remove(key)
    #         i = 0
    #         modified_dict = {}
    #         for key, value in checkpoint.items():
    #             if 'scaling_factor' in key: continue
    #             if 'num_batches_tracked' in key: continue
    #             if 'weight_integer' in key: continue
    #             if 'min' in key or 'max' in key: continue
    #             modified_key = model_key_list[i]
    #             modified_dict[modified_key] = value
    #             i += 1
    #         logging.info(model.load_state_dict(modified_dict, strict=False))
    #     else:
    #         logging.info("=> no checkpoint found at '{}'".format(args.resume))

    if args.regular != True:
        quantize_arch = quantize_arch_dict[args.arch]
        model = quantize_arch(model)

    # if args.arch == "q_resmlp_v4":
    #     for i in range(0, 24):
    #         model.blocks[i].inner.weight.requires_grad = False
    #         model.blocks[i].inner.bias.requires_grad = False
    #         model.blocks[i].outer.weight.requires_grad = False
    
    if args.freeze_w:
        for param in model.parameters():
            param.requires_grad = False

        for n, m in model.named_modules():
            if isinstance(m, LSQObserver):
                for param in m.parameters():
                    param.requires_grad = True

    print("args.batch_size", args.batch_size)
    # for name, m in model.named_modules():
    #     setattr(m, 'regular', args.regular)
    #     setattr(m, 'act_percentile', args.act_percentile)
    
    #logging.info(model)

    if args.resume:# and args.resume_quantize:
        if os.path.isfile(args.resume):
            logging.info("=> loading quantized checkpoint '{}'".format(args.resume))
            model.load_state_dict(torch.load(args.resume)['state_dict'])
        else:
            logging.info("=> no quantized checkpoint found at '{}'".format(args.resume))
    
    if args.rca:
        logging.info("=> Applying RCA on model")
        rca_for_resmlp(model)

    if args.gpu is not None:
        torch.cuda.set_device(args.gpu)
        model = model.cuda(args.gpu)
        if args.distill_method != 'None':
            teacher = teacher.cuda(args.gpu)
    else:
        # DataParallel will divide and allocate batch_size to all available GPUs
        if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
            model.features = torch.nn.DataParallel(model.features)
            model.cuda()
            # teacher is not alexnet or vgg
            if args.distill_method != 'None':
                teacher = torch.nn.DataParallel(teacher).cuda()
        else:
            model = torch.nn.DataParallel(model).cuda()
            if args.distill_method != 'None':
                teacher = torch.nn.DataParallel(teacher).cuda()

    # define loss function (criterion) and optimizer
    criterion = nn.CrossEntropyLoss().cuda(args.gpu)

    
    
    # enlarge weight decay for inner linear's bias
    if args.arch == "q_resmlp_v4":
        optimizer = torch.optim.SGD([
            {'params': (p for name, p in model.named_parameters() if "outer.weight" not in name), 'lr': args.lr},
            {'params': (p for name, p in model.named_parameters() if "outer.weight" in name), 'lr': args.lr*0.1},
        ], 
        # lr=args.lr,
        momentum=args.momentum,
        weight_decay=args.weight_decay)

    else:
        optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                    momentum=args.momentum,
                                    weight_decay=args.weight_decay)
    
    scheduler = CosineLRScheduler(optimizer, t_initial=args.epochs)
    cudnn.benchmark = True

    # Data loading code
    traindir = os.path.join(args.data, 'train')
    valdir = os.path.join(args.data, 'val')
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    train_resolution = 224
    if args.arch == "inceptionv3":
        train_resolution = 299

    train_dataset = datasets.ImageFolder(
        traindir,
        transforms.Compose([
            transforms.RandomResizedCrop(train_resolution),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]))

    if args.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
    else:
        train_sampler = None

    dataset_length = int(len(train_dataset) * args.data_percentage)
    if args.data_percentage == 1:
        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)
    else:
        partial_train_dataset, _ = torch.utils.data.random_split(train_dataset,
                                                                 [dataset_length, len(train_dataset) - dataset_length])
        train_loader = torch.utils.data.DataLoader(
            partial_train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
            num_workers=args.workers, pin_memory=True, sampler=train_sampler)

    test_resolution = (256, 224)
    if args.arch == 'inceptionv3':
        test_resolution = (342, 299)

    # evaluate on validation set
    val_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder(valdir, transforms.Compose([
            transforms.Resize(test_resolution[0]),
            transforms.CenterCrop(test_resolution[1]),
            transforms.ToTensor(),
            normalize,
        ])),
        batch_size=args.batch_size, shuffle=False,
        num_workers=args.workers, pin_memory=True)

    if args.evaluate:
        validate(val_loader, model, criterion, args)
        return

    # train code
    # wandb initialization
    if args.wandb:
        # wandb.login(key=)
        id = wandb.util.generate_id() #'33034pb2'
        wandb.init(
            project="resmlp_qat",
            id=id,
            resume=("must" if args.resume is True else False),
            config={
                "epochs": args.epochs,
                "learning_rate": args.lr,
                "batch_size": args.batch_size,
                "dataset": "imagenet",
                "data_percentage": args.data_percentage,
            })

    best_epoch = 0
    for epoch in range(args.start_epoch, args.epochs):
        # adjust_learning_rate(optimizer, epoch, args)

        # train for one epoch
        train(train_loader, model, criterion, optimizer, epoch, args)
        
        if not args.skip_val:
            acc1 = validate(val_loader, model, criterion, args)
            scheduler.step(epoch)

            # remember best acc@1 and save checkpoint
            is_best = acc1 > best_acc1
            best_acc1 = max(acc1, best_acc1)

            # record the best epoch
            if is_best:
                best_epoch = epoch
            
            logging.info(f'Best acc at epoch {best_epoch}: {best_acc1}')
            
            if not os.path.exists(args.save_path):
                os.makedirs(args.save_path)
            save_checkpoint({
                'epoch': epoch + 1,
                'arch': args.arch,
                'state_dict': model.state_dict(),
                'best_acc1': best_acc1,
                'optimizer': optimizer.state_dict(),
            }, is_best, args.save_path)
        else:
            # saving
            if not os.path.exists(args.save_path):
                os.makedirs(args.save_path)

            save_checkpoint({
                'epoch': epoch + 1,
                'arch': args.arch,
                'state_dict': model.state_dict(),
                'optimizer': optimizer.state_dict(),
            }, True, args.save_path)
        


def train(train_loader, model, criterion, optimizer, epoch, args):
    batch_time = AverageMeter('Time', ':6.3f')
    data_time = AverageMeter('Data', ':6.3f')
    losses = AverageMeter('Loss', ':.4e')
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    progress = ProgressMeter(
        len(train_loader),
        [batch_time, data_time, losses, top1, top5],
        prefix="Epoch: [{}]".format(epoch))

    # switch to train mode
    # set_training(model, True)
    model.train()
    end = time.time()
    for i, (images, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        if args.gpu is not None:
            images = images.cuda(args.gpu, non_blocking=True)
        target = target.cuda(args.gpu, non_blocking=True)

        # compute output
        output = model(images)
        loss = criterion(output, target)

        # measure accuracy and record loss
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        losses.update(loss.item(), images.size(0))
        top1.update(acc1[0], images.size(0))
        top5.update(acc5[0], images.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:
            progress.display(i)
            if args.wandb:
                to_log = {
                    "train/train_loss": loss.item(), 
                    "train/train_acc1": acc1[0], 
                    "train/train_acc5": acc5[0]
                }

                scales = model.get_scales()
                for scale in scales:
                    # to_log[f"train_quant/align_{i}"] = scale[0]
                    to_log[f"train_scales/{scale[0]}"] = scale[1].item()
                wandb.log(to_log)



# def train_kd(train_loader, model, teacher, criterion, optimizer, epoch, val_loader, args, ngpus_per_node,
#              dataset_length):
#     batch_time = AverageMeter('Time', ':6.3f')
#     data_time = AverageMeter('Data', ':6.3f')
#     losses = AverageMeter('Loss', ':.4e')
#     top1 = AverageMeter('Acc@1', ':6.2f')
#     top5 = AverageMeter('Acc@5', ':6.2f')
#     progress = ProgressMeter(
#         len(train_loader),
#         [batch_time, data_time, losses, top1, top5],
#         prefix="Epoch: [{}]".format(epoch))

#     # switch to train mode
#     if args.fix_BN == True:
#         model.eval()
#     else:
#         model.train()
#     teacher.eval()

#     end = time.time()

#     for i, (images, target) in enumerate(train_loader):
#         # measure data loading time
#         data_time.update(time.time() - end)

#         if args.gpu is not None:
#             images = images.cuda(args.gpu, non_blocking=True)
#         target = target.cuda(args.gpu, non_blocking=True)

#         # compute output
#         output = model(images)
#         if args.distill_method != 'None':
#             with torch.no_grad():
#                 teacher_output = teacher(images)

#         if args.distill_method == 'None':
#             loss = criterion(output, target)
#         elif args.distill_method == 'KD_naive':
#             loss = loss_kd(output, target, teacher_output, args)
#         else:
#             raise NotImplementedError

#         # measure accuracy and record loss
#         acc1, acc5 = accuracy(output, target, topk=(1, 5))
#         losses.update(loss.item(), images.size(0))
#         top1.update(acc1[0], images.size(0))
#         top5.update(acc5[0], images.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:
#             progress.display(i)
#         if i % args.print_freq == 0 and args.rank == 0:
#             print('Epoch {epoch_} [{iters}]  Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(epoch_=epoch, iters=i,
#                                                                                                top1=top1, top5=top5))
#             if args.wandb:
#                 to_log = {
#                     "train_loss": loss.item(), 
#                     "train_acc1": acc1[0], 
#                     "train_acc5": acc5[0]
#                 }

#                 scales = model.get_scales()
#                 # for i, scale in enumerate(scales):
#                 #     to_log[f"train_quant/scale_{i}"] = scale
#                 # wandb.log(to_log)

#         if i % ((dataset_length // (
#                 args.batch_size * args.evaluate_times)) + 2) == 0 and i > 0 and args.evaluate_times > 0:
#             acc1 = validate(val_loader, model, criterion, args)

#             # switch to train mode
#             if args.fix_BN == True:
#                 model.eval()
#             else:
#                 model.train()

#             # remember best acc@1 and save checkpoint
#             global best_acc1
#             is_best = acc1 > best_acc1
#             best_acc1 = max(acc1, best_acc1)

#             if not args.multiprocessing_distributed or (args.multiprocessing_distributed
#                                                         and args.rank % ngpus_per_node == 0):
#                 if not os.path.exists(args.save_path):
#                     os.makedirs(args.save_path)

#                 save_checkpoint({
#                     'epoch': epoch + 1,
#                     'arch': args.arch,
#                     'state_dict': model.state_dict(),
#                     'best_acc1': best_acc1,
#                     'optimizer': optimizer.state_dict(),
#                 }, is_best, args.save_path)
#                 # print(model.state_dict())
#                 print("Saved checkpoint.")


def validate(val_loader, model, criterion, args):
    batch_time = AverageMeter('Time', ':6.3f')
    losses = AverageMeter('Loss', ':.4e')
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    progress = ProgressMeter(
        len(val_loader),
        [batch_time, losses, top1, top5],
        prefix='Test: ')

    # switch to evaluate mode
    model.eval()

    with torch.no_grad():
        end = time.time()
        for i, (images, target) in enumerate(val_loader):
            if args.gpu is not None:
                images = images.cuda(args.gpu, non_blocking=True)
            target = target.cuda(args.gpu, non_blocking=True)

            # compute output
            output = model(images)
            loss = criterion(output, target)

            # measure accuracy and record loss
            acc1, acc5 = accuracy(output, target, topk=(1, 5))
            losses.update(loss.item(), images.size(0))
            top1.update(acc1[0], images.size(0))
            top5.update(acc5[0], images.size(0))

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if i % args.print_freq == 0:
                progress.display(i)

        logging.info(' * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(top1=top1, top5=top5))

    if args.wandb:
        to_log = {
            "evaluation/test_avg_loss": losses.avg, 
            "evaluation/test_avg_acc1": top1.avg, 
            "evaluation/test_avg_acc5": top5.avg,
        }
        
        # scales = model.get_scales()
        # for i, scale in enumerate(scales):
        #     to_log[f"eval_quant/scale_{i}"] = scale
        
        wandb.log(to_log)

    torch.save({'observer.scale': {k: v for k, v in model.state_dict().items() if 'observer.scale' in k},
                'w_int': {k: v for k, v in model.state_dict().items() if 'w_int' in k},
                'b_int': {k: v for k, v in model.state_dict().items() if 'b_int' in k},
                'mult': {k: v for k, v in model.state_dict().items() if 'mult' in k},
                'shift': {k: v for k, v in model.state_dict().items() if 'shift' in k},
                'res_mult': {k: v for k, v in model.state_dict().items() if 'res_mult' in k},
                'res_shift': {k: v for k, v in model.state_dict().items() if 'res_shift' in k},
                }, args.save_path + 'quantized_checkpoint.pth.tar')

                
    # logging.info(model.state_dict().items())
    return top1.avg


def save_checkpoint(state, is_best, filename=None):
    torch.save(state, filename + 'checkpoint.pth.tar')
    if is_best:
        shutil.copyfile(filename + 'checkpoint.pth.tar', filename + 'model_best.pth.tar')


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].reshape(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res


def loss_kd(output, target, teacher_output, args):
    """
    Compute the knowledge-distillation (KD) loss given outputs and labels.
    "Hyperparameters": temperature and alpha
    The KL Divergence for PyTorch comparing the softmaxs of teacher and student.
    The KL Divergence expects the input tensor to be log probabilities.
    """
    alpha = args.distill_alpha
    T = args.temperature
    KD_loss = F.kl_div(F.log_softmax(output / T, dim=1), F.softmax(teacher_output / T, dim=1)) * (alpha * T * T) + \
              F.cross_entropy(output, target) * (1. - alpha)

    return KD_loss


if __name__ == '__main__':
    main()