infer.py

# Copyright (c) 2015-present, Facebook, Inc.
# All rights reserved.
import argparse
import datetime
import numpy as np
import time
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import json

import tqdm
import lightning as L

from pathlib import Path

from timm.data import Mixup
from timm.models import create_model
from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
from timm.scheduler import create_scheduler
from timm.optim import create_optimizer
from timm.utils import NativeScaler, get_state_dict, ModelEma

from datasets import build_dataset
from engine import train_one_epoch, evaluate
from losses import DistillationLoss
from samplers import RASampler
from functools import partial


from vit import VisionTransformerDiffPruning
from lvvit import LVViTDiffPruning


def get_args_parser():
    parser = argparse.ArgumentParser('DeiT training and evaluation script', add_help=False)
    parser.add_argument('--batch-size', default=128, type=int)
    parser.add_argument('--arch', default='deit_small', type=str, help='Name of model to train')
    parser.add_argument('--input-size', default=224, type=int, help='images input size')
    parser.add_argument('--data-path', default='/datasets01/imagenet_full_size/061417/', type=str,
                        help='dataset path')
    parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19'],
                        type=str, help='Image Net dataset path')
    parser.add_argument('--inat-category', default='name',
                        choices=['kingdom', 'phylum', 'class', 'order', 'supercategory', 'family', 'genus', 'name'],
                        type=str, help='semantic granularity')
    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--model-path', default='', help='resume from checkpoint')
    parser.add_argument('--num_workers', default=10, type=int)
    parser.add_argument('--pin-mem', action='store_true',
                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
    parser.add_argument('--no-pin-mem', action='store_false', dest='pin_mem',
                        help='')
    parser.set_defaults(pin_mem=True)
    parser.add_argument('--base_rate', type=float, default=0.7)

    return parser


def main(args):

    cudnn.benchmark = True
    dataset_val, _ = build_dataset(is_train=False, args=args)

    data_loader_val = torch.utils.data.DataLoader(
        dataset_val,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=args.pin_mem,
        drop_last=False
    )

    base_rate = args.base_rate
    KEEP_RATE = [base_rate, base_rate ** 2, base_rate ** 3]

    if args.arch == 'deit_small':
        PRUNING_LOC = [3,6,9] 
        print(f"Creating model: {args.arch}")
        print('token_ratio =', KEEP_RATE, 'at layer', PRUNING_LOC)
        model = VisionTransformerDiffPruning(
            patch_size=16, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4, qkv_bias=True, 
            pruning_loc=PRUNING_LOC, token_ratio=KEEP_RATE
            )
    elif args.arch == 'deit_256':
        PRUNING_LOC = [3,6,9] 
        print(f"Creating model: {args.arch}")
        print('token_ratio =', KEEP_RATE, 'at layer', PRUNING_LOC)
        model = VisionTransformerDiffPruning(
            patch_size=16, embed_dim=256, depth=12, num_heads=4, mlp_ratio=4, qkv_bias=True, 
            pruning_loc=PRUNING_LOC, token_ratio=KEEP_RATE
            )
    elif args.arch == 'lvvit_s':
        PRUNING_LOC = [4,8,12] 
        print(f"Creating model: {args.arch}")
        print('token_ratio =', KEEP_RATE, 'at layer', PRUNING_LOC)
        model = LVViTDiffPruning(
            patch_size=16, embed_dim=384, depth=16, num_heads=6, mlp_ratio=3.,
            p_emb='4_2',skip_lam=2., return_dense=True,mix_token=True,
            pruning_loc=PRUNING_LOC, token_ratio=KEEP_RATE
        )
    elif args.arch == 'lvvit_m':
        PRUNING_LOC = [5,10,15] 
        print(f"Creating model: {args.arch}")
        print('token_ratio =', KEEP_RATE, 'at layer', PRUNING_LOC)
        model = LVViTDiffPruning(
            patch_size=16, embed_dim=512, depth=20, num_heads=8, mlp_ratio=3.,
            p_emb='4_2',skip_lam=2., return_dense=True,mix_token=True,
            pruning_loc=PRUNING_LOC, token_ratio=KEEP_RATE
        )
    else:
        raise NotImplementedError

    model_path = args.model_path
    checkpoint = torch.load(model_path, map_location="cpu")
    model.load_state_dict(checkpoint["model"])

    print('## model has been successfully loaded')

    model = model.cuda()

    n_parameters = sum(p.numel() for p in model.parameters())
    print('number of params:', n_parameters)

    criterion = torch.nn.CrossEntropyLoss().cuda()
    validate(data_loader_val, model, criterion)

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self, name, fmt=':f'):
        self.name = name
        self.fmt = fmt
        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 __str__(self):
        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
        return fmtstr.format(**self.__dict__)

class ProgressMeter(object):
    def __init__(self, num_batches, meters, prefix=""):
        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
        self.meters = meters
        self.prefix = prefix

    def display(self, batch):
        entries = [self.prefix + self.batch_fmtstr.format(batch)]
        entries += [str(meter) for meter in self.meters]
        print('\t'.join(entries))

    def _get_batch_fmtstr(self, num_batches):
        num_digits = len(str(num_batches // 1))
        fmt = '{:' + str(num_digits) + 'd}'
        return '[' + fmt + '/' + fmt.format(num_batches) + ']'

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 validate(args, val_loader, model, criterion):
    batch_time = AverageMeter('Time (ms)', ':.3f')
    losses = AverageMeter('Loss', ':.4e')
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    model.eval()
    
    ### Create fabric instance
    fabric = L.Fabric(
        accelerator='cuda',
        strategy='dp',
        devices=1,
        num_nodes=1,
        precision='32',
    )
    fabric.launch()

    ### Wrap model with fabric
    model       = fabric.setup_module(model, move_to_device=True)
    val_loader  = fabric.setup_dataloaders(val_loader, use_distributed_sampler=False, move_to_device=True)
    
    ### Use TQDM instead
    tqdm_progress_bar = tqdm(val_loader)

    ### Create torch.profiler instance
    ### If we are using tensorboard, wrap evaluation calls with it
    if args.eval_tensorboard:
        torchprofiler = torch.profiler.profile(
            ### Create profiler instance
            activities=[torch.profiler.ProfilerActivity.CUDA, torch.profiler.ProfilerActivity.CPU],
            schedule=torch.profiler.schedule(wait=24, warmup=25, active=1),
            on_trace_ready=torch.profiler.tensorboard_trace_handler("./bin/"),
            record_shapes=True,
            with_stack=True,
            with_flops=True,
            with_modules=True,
        )

        ### Start profiling!
        torchprofiler.start()
    else:
        torchprofiler = None

    ### Wait for profiler to be setup
    time.sleep(5.0)

    with torch.no_grad():
        start_event             = torch.cuda.Event(enable_timing=True)
        end_event               = torch.cuda.Event(enable_timing=True)

        for batch_index, (images, target) in enumerate(tqdm_progress_bar):
            ### Abort early 
            if args.forward_pass_count is not None and batch_index > args.forward_pass_count:
                print('infer.py: Exiting Early')
                break

            ### Simple Warmup 
            if batch_index < 1:
                for k in range(25):
                    _ = model(images)
                torch.cuda.synchronize()

            ### Start recording
            start_event.record()

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

            end_event.record()
            torch.cuda.synchronize()

            # 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( start_event.elapsed_time( end_event ) )

            tqdm_progress_bar.set_description(
                desc='Acc@1 Avg: {:.2f} | Average Batch Time (ms): {:.2f}'.format(top1.avg, batch_time.avg),
                refresh=True
            )

            if torchprofiler is not None:
                torchprofiler.step()

        ###
        ### Print info
        ###
        print(' * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(top1=top1, top5=top5))
        print('Average Batch Time: {:.3f}'.format(batch_time.avg))
        print('Average Loss: {:.3f}'.format(losses.avg))

    return top1.avg

if __name__ == '__main__':
    parser = argparse.ArgumentParser('Dynamic evaluation script', parents=[get_args_parser()])
    args = parser.parse_args()
    main(args)