test.py

import argparse
import time
import yaml
import json
import os
import numpy as np
import logging
from collections import OrderedDict
from contextlib import suppress
from datetime import datetime
from spikingjelly.clock_driven import functional
from spikingjelly.datasets.cifar10_dvs import CIFAR10DVS
from spikingjelly.datasets.n_caltech101 import NCaltech101

import torch
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as NativeDDP
from torchvision import transforms

from timm.data import create_dataset, create_loader, resolve_data_config
from timm.models import (
    create_model,
    safe_model_name,
    resume_checkpoint,
    load_checkpoint,
    convert_splitbn_model,
)
from timm.utils import *
from timm.utils import ApexScaler, NativeScaler
import dvs_utils

try:
    from apex import amp
    from apex.parallel import DistributedDataParallel as ApexDDP
    from apex.parallel import convert_syncbn_model

    has_apex = True
except ImportError:
    has_apex = False

has_native_amp = False
try:
    if getattr(torch.cuda.amp, "autocast") is not None:
        has_native_amp = True
except AttributeError:
    pass

try:
    import wandb

    has_wandb = True
except ImportError:
    has_wandb = False

torch.backends.cudnn.benchmark = True
# The first arg parser parses out only the --config argument, this argument is used to
# load a yaml file containing key-values that override the defaults for the main parser below
config_parser = parser = argparse.ArgumentParser(
    description="Training Config", add_help=False
)
parser.add_argument(
    "-c",
    "--config",
    default="imagenet.yml",
    type=str,
    metavar="FILE",
    help="YAML config file specifying default arguments",
)

parser = argparse.ArgumentParser(description="PyTorch ImageNet Training")

# Dataset / Model parameters
parser.add_argument("-data-dir", metavar="DIR", default="", help="path to dataset")
parser.add_argument("--dataset", "-d", metavar="NAME", default="torch/cifar10", help="dataset type (default: ImageFolder/ImageTar if empty)")
parser.add_argument("--train-split", metavar="NAME", default="train", help="dataset train split (default: train)")
parser.add_argument("--val-split", metavar="NAME", default="validation", help="dataset validation split (default: validation)")
parser.add_argument("--model", default="spikeformer", type=str, metavar="MODEL", help='Name of model to train (default: "countception")')
parser.add_argument("--pooling_stat", default="1111", type=str, help="pooling layers in embedding layer")
parser.add_argument("--TET", default=False, type=bool)
parser.add_argument("--TET-means", default=1.0, type=float)
parser.add_argument("--TET-lamb", default=0.0, type=float)
parser.add_argument("--spike-mode", default="lif", type=str)
parser.add_argument("--layer", default=4, type=int)
parser.add_argument("--in-channels", default=3, type=int)
parser.add_argument("--pretrained", action="store_true", default=False, help="Start with pretrained version of specified network (if avail)")
parser.add_argument("--initial-checkpoint", default="", type=str, metavar="PATH", help="Initialize model from this checkpoint (default: none)")
parser.add_argument("--resume", default="", type=str, metavar="PATH", help="Resume full model and optimizer state from checkpoint (default: none)")
parser.add_argument("--no-resume-opt", action="store_true", default=False, help="prevent resume of optimizer state when resuming model")
parser.add_argument("--num-classes", type=int, default=1000, metavar="N", help="number of label classes (Model default if None)")
parser.add_argument("--time-steps", type=int, default=4, metavar="N")
parser.add_argument("--num-heads", type=int, default=8, metavar="N")
parser.add_argument("--patch-size", type=int, default=None, metavar="N", help="Image patch size")
parser.add_argument("--mlp-ratio", type=int, default=4, metavar="N", help="expand ration of embedding dimension in MLP block")
parser.add_argument("--gp", default=None, type=str, metavar="POOL", help="Global pool type, one of (fast, avg, max, avgmax, avgmaxc). Model default if None.")
parser.add_argument("--img-size", type=int, default=None, metavar="N", help="Image patch size (default: None => model default)")
parser.add_argument("--input-size", default=None, nargs=3, type=int, metavar="N N N", help="Input all image dimensions (d h w, e.g. --input-size 3 224 224), uses model default if empty")
parser.add_argument("--crop-pct", default=None, type=float, metavar="N", help="Input image center crop percent (for validation only)")
parser.add_argument("--mean", type=float, nargs="+", default=None, metavar="MEAN", help="Override mean pixel value of dataset")
parser.add_argument("--std", type=float, nargs="+", default=None, metavar="STD", help="Override std deviation of of dataset" )
parser.add_argument("--interpolation", default="", type=str,metavar="NAME", help="Image resize interpolation type (overrides model)")
parser.add_argument("-b", "--batch-size", type=int, default=32, metavar="N", help="input batch size for training (default: 32)")
parser.add_argument("-vb", "--val-batch-size", type=int, default=16, metavar="N", help="input val batch size for training (default: 32)")

# Augmentation & regularization parameters
parser.add_argument("--no-aug", action="store_true", default=False, help="Disable all training augmentation, override other train aug args")
parser.add_argument("--scale", type=float, nargs="+", default=[0.08, 1.0], metavar="PCT", help="Random resize scale (default: 0.08 1.0)")
parser.add_argument("--ratio", type=float, nargs="+", default=[3.0 / 4.0, 4.0 / 3.0], metavar="RATIO", help="Random resize aspect ratio (default: 0.75 1.33)")
parser.add_argument("--hflip", type=float, default=0.5, help="Horizontal flip training aug probability")
parser.add_argument("--vflip", type=float, default=0.0, help="Vertical flip training aug probability")
parser.add_argument("--color-jitter", type=float, default=0.4, metavar="PCT", help="Color jitter factor (default: 0.4)")
parser.add_argument("--aa", type=str, default=None, metavar="NAME", help='Use AutoAugment policy. "v0" or "original". (default: None)')
parser.add_argument("--aug-splits", type=int, default=0, help="Number of augmentation splits (default: 0, valid: 0 or >=2)")
parser.add_argument("--jsd", action="store_true", default=False, help="Enable Jensen-Shannon Divergence + CE loss. Use with `--aug-splits`.")
parser.add_argument("--bce-loss", action="store_true", default=False, help="Enable BCE loss w/ Mixup/CutMix use.")
parser.add_argument("--bce-target-thresh", type=float, default=None, help="Threshold for binarizing softened BCE targets (default: None, disabled)")
parser.add_argument("--reprob", type=float, default=0.0, metavar="PCT", help="Random erase prob (default: 0.)")
parser.add_argument("--remode", type=str, default="const", help='Random erase mode (default: "const")')
parser.add_argument("--recount", type=int, default=1, help="Random erase count (default: 1)")
parser.add_argument("--resplit", action="store_true", default=False, help="Do not random erase first (clean) augmentation split")
parser.add_argument("--mixup", type=float, default=0.0, help="mixup alpha, mixup enabled if > 0. (default: 0.)")
parser.add_argument("--cutmix", type=float, default=0.0, help="cutmix alpha, cutmix enabled if > 0. (default: 0.)")
parser.add_argument("--cutmix-minmax", type=float, nargs="+", default=None, help="cutmix min/max ratio, overrides alpha and enables cutmix if set (default: None)")
parser.add_argument("--mixup-prob", type=float, default=1.0, help="Probability of performing mixup or cutmix when either/both is enabled")
parser.add_argument("--mixup-switch-prob", type=float, default=0.5, help="Probability of switching to cutmix when both mixup and cutmix enabled")
parser.add_argument("--mixup-mode", type=str, default="batch", help='How to apply mixup/cutmix params. Per "batch", "pair", or "elem"')
parser.add_argument("--mixup-off-epoch", default=0, type=int, metavar="N", help="Turn off mixup after this epoch, disabled if 0 (default: 0)")
parser.add_argument("--smoothing", type=float, default=0.1, help="Label smoothing (default: 0.1)")
parser.add_argument("--train-interpolation", type=str, default="random", help='Training interpolation (random, bilinear, bicubic default: "random")')
parser.add_argument("--drop", type=float, default=0.0, metavar="PCT", help="Dropout rate (default: 0.)")
parser.add_argument("--drop-connect", type=float, default=None, metavar="PCT", help="Drop connect rate, DEPRECATED, use drop-path (default: None)")
parser.add_argument("--drop-path", type=float, default=0.2, metavar="PCT", help="Drop path rate (default: None)")
parser.add_argument( "--drop-block", type=float, default=None, metavar="PCT", help="Drop block rate (default: None)")

# Batch norm parameters (only works with gen_efficientnet based models currently)
parser.add_argument("--bn-tf", action="store_true", default=False, help="Use Tensorflow BatchNorm defaults for models that support it (default: False)")
parser.add_argument("--bn-momentum", type=float, default=None, help="BatchNorm momentum override (if not None)")
parser.add_argument("--bn-eps", type=float, default=None, help="BatchNorm epsilon override (if not None)")
parser.add_argument("--sync-bn", action="store_true", help="Enable NVIDIA Apex or Torch synchronized BatchNorm.")
parser.add_argument("--dist-bn", type=str, default="", help='Distribute BatchNorm stats between nodes after each epoch ("broadcast", "reduce", or "")')
parser.add_argument("--split-bn", action="store_true", help="Enable separate BN layers per augmentation split.")

# Model Exponential Moving Average
parser.add_argument("--model-ema", action="store_true", default=False, help="Enable tracking moving average of model weights")
parser.add_argument("--model-ema-force-cpu", action="store_true", default=False, help="Force ema to be tracked on CPU, rank=0 node only. Disables EMA validation.")
parser.add_argument("--model-ema-decay", type=float, default=0.9998, help="decay factor for model weights moving average (default: 0.9998)")

# Misc
parser.add_argument("--seed", type=int, default=42, metavar="S", help="random seed (default: 42)")
parser.add_argument("--log-interval", type=int, default=100, metavar="N", help="how many batches to wait before logging training status")
parser.add_argument("--recovery-interval", type=int, default=0, metavar="N", help="how many batches to wait before writing recovery checkpoint")
parser.add_argument("--checkpoint-hist", type=int, default=10, metavar="N", help="number of checkpoints to keep (default: 10)")
parser.add_argument("-j"  "--workers", type=int, default=4, metavar="N", help="how many training processes to use (default: 1)")
parser.add_argument("--save-images", action="store_true", default=False, help="save images of input bathes every log interval for debugging")
parser.add_argument("--amp", action="store_true", default=False, help="use NVIDIA Apex AMP or Native AMP for mixed precision training")
parser.add_argument("--apex-amp", action="store_true", default=False, help="Use NVIDIA Apex AMP mixed precision")
parser.add_argument("--native-amp", action="store_true", default=False, help="Use Native Torch AMP mixed precision")
parser.add_argument("--channels-last", action="store_true", default=False, help="Use channels_last memory layout")
parser.add_argument("--pin-mem", action="store_true", default=False, help="Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.")
parser.add_argument("--no-prefetcher", action="store_true", default=False, help="disable fast prefetcher")
parser.add_argument("--dvs-aug", action="store_true", default=False, help="disable fast prefetcher")
parser.add_argument("--dvs-trival-aug", action="store_true", default=False, help="disable fast prefetcher")
parser.add_argument("--save-qkv", action="store_true", default=False, help="disable fast prefetcher")
parser.add_argument("--output", default="", type=str, metavar="PATH", help="path to output folder (default: none, current dir)")
parser.add_argument("--experiment", default="", type=str, metavar="NAME", help="name of train experiment, name of sub-folder for output")
parser.add_argument("--eval-metric", default="top1", type=str, metavar="EVAL_METRIC", help='Best metric (default: "top1")')
parser.add_argument("--tta", type=int, default=0, metavar="N", help="Test/inference time augmentation (oversampling) factor. 0=None (default: 0)")
parser.add_argument("--local_rank", default=0, type=int)
parser.add_argument("--use-multi-epochs-loader", action="store_true", default=False, help="use the multi-epochs-loader to save time at the beginning of every epoch")
parser.add_argument("--large-valid", action="store_true", default=False, help="use the multi-epochs-loader to save time at the beginning of every epoch")
parser.add_argument("--torchscript", dest="torchscript", action="store_true", help="convert model torchscript for inference")
parser.add_argument("--log-wandb", action="store_true", default=False, help="log training and validation metrics to wandb")
parser.add_argument("--attention_mode", type=str, default="T_STAtten", help='choose one of methods: J_STAtten, T_STAtten, N_STAtten, SDT')

_logger = logging.getLogger("valid")
stream_handler = logging.StreamHandler()
format_str = "%(asctime)s %(levelname)s: %(message)s"
stream_handler.setFormatter(logging.Formatter(format_str))
_logger.addHandler(stream_handler)
_logger.propagate = False


def _parse_args():
    # Do we have a config file to parse?
    args_config, remaining = config_parser.parse_known_args()
    if args_config.config:
        with open(args_config.config, "r") as f:
            cfg = yaml.safe_load(f)
            parser.set_defaults(**cfg)

    # The main arg parser parses the rest of the args, the usual
    # defaults will have been overridden if config file specified.
    args = parser.parse_args(remaining)

    # Cache the args as a text string to save them in the output dir later
    args_text = yaml.safe_dump(args.__dict__, default_flow_style=False)
    return args, args_text


def main():
    setup_default_logging()
    args, args_text = _parse_args()

    args.prefetcher = not args.no_prefetcher
    args.distributed = False
    if "WORLD_SIZE" in os.environ:
        args.distributed = int(os.environ["WORLD_SIZE"]) > 1
    args.device = "cuda:1"
    args.world_size = 1
    args.rank = 0  # global rank
    if args.distributed:
        args.device = "cuda:%d" % args.local_rank
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(backend="nccl", init_method="env://")
        args.world_size = torch.distributed.get_world_size()
        args.rank = torch.distributed.get_rank()
        _logger.info(
            "Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d."
            % (args.rank, args.world_size)
        )
    else:
        _logger.info("Training with a single process on 1 GPUs.")
    assert args.rank >= 0

    # resolve AMP arguments based on PyTorch / Apex availability
    use_amp = None
    if args.amp:
        # `--amp` chooses native amp before apex (APEX ver not actively maintained)
        if has_native_amp:
            args.native_amp = True
        elif has_apex:
            args.apex_amp = True
    if args.apex_amp and has_apex:
        use_amp = "apex"
    elif args.native_amp and has_native_amp:
        use_amp = "native"
    elif args.apex_amp or args.native_amp:
        _logger.warning(
            "Neither APEX or native Torch AMP is available, using float32. "
            "Install NVIDA apex or upgrade to PyTorch 1.6"
        )

    torch.backends.cudnn.benchmark = True
    os.environ["PYTHONHASHSEED"] = str(args.seed)
    np.random.seed(args.seed)
    torch.initial_seed()  # dataloader multi processing
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    random_seed(args.seed, args.rank)

    args.dvs_mode = False
    if args.dataset in ["cifar10-dvs", "ncaltech101"]:
        args.dvs_mode = True

    model = create_model(
        args.model,
        T=args.time_steps,
        pretrained=args.pretrained,
        drop_rate=args.drop,
        drop_path_rate=args.drop_path,
        drop_block_rate=args.drop_block,
        num_heads=args.num_heads,
        num_classes=args.num_classes,
        pooling_stat=args.pooling_stat,
        img_size_h=args.img_size,
        img_size_w=args.img_size,
        patch_size=args.patch_size,
        embed_dims=args.dim,
        mlp_ratios=args.mlp_ratio,
        in_channels=args.in_channels,
        qkv_bias=False,
        depths=args.layer,
        sr_ratios=1,
        spike_mode=args.spike_mode,
        dvs_mode=args.dvs_mode,
        TET=args.TET,
    )
    if args.local_rank == 0:
        _logger.info("Creating model")
        n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
        _logger.info(f"number of params: {n_parameters}")

    if args.num_classes is None:
        assert hasattr(
            model, "num_classes"
        ), "Model must have `num_classes` attr if not set on cmd line/config."
        args.num_classes = (
            model.num_classes
        )  # FIXME handle model default vs config num_classes more elegantly

    if args.local_rank == 0:
        _logger.info(
            f"Model {safe_model_name(args.model)} created, param count:{sum([m.numel() for m in model.parameters()])}"
        )

    data_config = resolve_data_config(
        vars(args), model=model, verbose=args.local_rank == 0
    )

    # setup augmentation batch splits for contrastive loss or split bn
    num_aug_splits = 0
    if args.aug_splits > 0:
        assert args.aug_splits > 1, "A split of 1 makes no sense"
        num_aug_splits = args.aug_splits

    # enable split bn (separate bn stats per batch-portion)
    if args.split_bn:
        assert num_aug_splits > 1 or args.resplit
        model = convert_splitbn_model(model, max(num_aug_splits, 2))

    # move model to GPU, enable channels last layout if set
    model.cuda()
    if args.channels_last:
        model = model.to(memory_format=torch.channels_last)

    # setup synchronized BatchNorm for distributed training
    if args.distributed and args.sync_bn:
        assert not args.split_bn
        if has_apex and use_amp != "native":
            # Apex SyncBN preferred unless native amp is activated
            model = convert_syncbn_model(model)
        else:
            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
        if args.local_rank == 0:
            _logger.info(
                "Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using "
                "zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled."
            )

    if args.torchscript:
        assert not use_amp == "apex", "Cannot use APEX AMP with torchscripted model"
        assert not args.sync_bn, "Cannot use SyncBatchNorm with torchscripted model"
        model = torch.jit.script(model)

    # setup automatic mixed-precision (AMP) loss scaling and op casting
    amp_autocast = suppress  # do nothing
    loss_scaler = None
    if use_amp == "apex":
        model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
        loss_scaler = ApexScaler()
        if args.local_rank == 0:
            _logger.info("Using NVIDIA APEX AMP. Training in mixed precision.")
    elif use_amp == "native":
        amp_autocast = torch.cuda.amp.autocast
        loss_scaler = NativeScaler()
        if args.local_rank == 0:
            _logger.info("Using native Torch AMP. Training in mixed precision.")
    else:
        if args.local_rank == 0:
            _logger.info("AMP not enabled. Training in float32.")

    # optionally resume from a checkpoint
    if args.resume:
        resume_checkpoint(
            model,
            args.resume,
            optimizer=None if args.no_resume_opt else optimizer,
            loss_scaler=None if args.no_resume_opt else loss_scaler,
            log_info=args.local_rank == 0,
        )

    # setup exponential moving average of model weights, SWA could be used here too
    model_ema = None
    if args.model_ema:
        # Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
        model_ema = ModelEmaV2(
            model,
            decay=args.model_ema_decay,
            device="cpu" if args.model_ema_force_cpu else None,
        )
        if args.resume:
            load_checkpoint(model_ema.module, args.resume, use_ema=True)

    if args.distributed:
        if has_apex and use_amp != "native":
            # Apex DDP preferred unless native amp is activated
            if args.local_rank == 0:
                _logger.info("Using NVIDIA APEX DistributedDataParallel.")
            model = ApexDDP(model, delay_allreduce=True, find_unused_parameters=True)
        else:
            if args.local_rank == 0:
                _logger.info("Using native Torch DistributedDataParallel.")
            model = NativeDDP(
                model, device_ids=[args.local_rank], find_unused_parameters=True
            )  # can use device str in Torch >= 1.1
        # NOTE: EMA model does not need to be wrapped by DDP

    # create the train and eval datasets
    dataset_eval = None, None
    if args.dataset == "cifar10-dvs":
        dataset = CIFAR10DVS(
            args.data_dir,
            data_type="frame",
            frames_number=args.time_steps,
            split_by="number",
        )
        _, dataset_eval = dvs_utils.split_to_train_test_set(0.9, dataset, 10)

    elif args.dataset == "ncaltech101":
        dataset = NCaltech101(
            args.data_dir,
            data_type="frame",
            frames_number=args.time_steps,
            split_by="number",
        )
        dataset_train, dataset_eval = dvs_utils.build_ncaltech(args.data_dir, True)
    else:
        dataset_eval = create_dataset(
            args.dataset,
            # root=args.data_dir,
            root="/gpfs/gibbs/project/panda/shared/imagenet/",
            split=args.val_split,
            is_training=False,
            batch_size=args.batch_size,
            download=True,
        )

    loader_eval = None
    if args.dataset in dvs_utils.DVS_DATASET:
        loader_eval = torch.utils.data.DataLoader(
            dataset_eval,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.workers,
            pin_memory=True,
        )
    elif args.dataset == "imagenet" and args.large_valid:
        dataset_eval.transform = transforms.Compose(
            [
                transforms.Resize(320),
                transforms.CenterCrop(288),
                transforms.ToTensor(),
                transforms.Normalize(mean=data_config["mean"], std=data_config["std"]),
            ]
        )
        sampler = torch.utils.data.distributed.DistributedSampler(dataset_eval)
        loader_eval = torch.utils.data.DataLoader(
            dataset_eval,
            batch_size=args.val_batch_size,
            num_workers=args.workers,
            sampler=sampler,
            pin_memory=args.pin_mem,
            drop_last=False,
        )
    else:
        loader_eval = create_loader(
            dataset_eval,
            input_size=data_config["input_size"],
            batch_size=args.val_batch_size,
            is_training=False,
            use_prefetcher=args.prefetcher,
            interpolation=data_config["interpolation"],
            mean=data_config["mean"],
            std=data_config["std"],
            num_workers=args.workers,
            distributed=args.distributed,
            crop_pct=data_config["crop_pct"],
            pin_memory=args.pin_mem,
        )

    validate_loss_fn = nn.CrossEntropyLoss().cuda()

    # setup checkpoint saver and eval metric tracking
    if args.experiment:
        exp_name = args.experiment
    else:
        exp_name = "-".join(
            [
                datetime.now().strftime("%Y%m%d-%H%M%S"),
                safe_model_name(args.model),
                "data-" + args.dataset.split("/")[-1],
                f"t-{args.time_steps}",
                f"spike-{args.spike_mode}",
            ]
        )
    output_dir = get_outdir(args.output if args.output else "./output/valid", exp_name)
    if args.rank == 0:
        file_handler = logging.FileHandler(
            os.path.join(output_dir, f"{args.model}.log"), "w"
        )
        file_handler.setFormatter(logging.Formatter(format_str))
        file_handler.setLevel(logging.INFO)
        _logger.addHandler(file_handler)

    try:
        if args.distributed and args.dist_bn in ("broadcast", "reduce"):
            if args.local_rank == 0:
                _logger.info("Distributing BatchNorm running means and vars")
            distribute_bn(model, args.world_size, args.dist_bn == "reduce")
        eval_metrics = validate(
            model,
            loader_eval,
            validate_loss_fn,
            args,
            output_dir=output_dir,
            amp_autocast=amp_autocast,
        )
        if args.local_rank == 0:
            non_zero_str = json.dumps(eval_metrics["non_zero"], indent=4)
            firing_rate_str = json.dumps(eval_metrics["firing_rate"], indent=4)
            _logger.info("top-1:", eval_metrics["top1"])
            _logger.info("non_zero: ")
            _logger.info(non_zero_str)
            _logger.info("firing_rate: ")
            _logger.info(firing_rate_str)
        if model_ema is not None and not args.model_ema_force_cpu:
            if args.distributed and args.dist_bn in ("broadcast", "reduce"):
                distribute_bn(model_ema, args.world_size, args.dist_bn == "reduce")
    except KeyboardInterrupt:
        pass


def validate(
    model, loader, loss_fn, args, output_dir=None, amp_autocast=suppress, log_suffix=""
):
    batch_time_m = AverageMeter()
    losses_m = AverageMeter()
    top1_m = AverageMeter()
    top5_m = AverageMeter()

    def calc_non_zero_rate(s_dict, nz_dict, idx, t):
        for k, v_ in s_dict.items():
            v = v_[t, ...]
            x_shape = torch.tensor(list(v.shape))
            all_neural = torch.prod(x_shape)
            z = torch.nonzero(v)
            if k in nz_dict.keys():
                nz_dict[k] += (z.shape[0] / all_neural).item() / idx
            else:
                nz_dict[k] = (z.shape[0] / all_neural).item() / idx
        return nz_dict

    def calc_firing_rate(s_dict, fr_dict, idx, t):
        for k, v_ in s_dict.items():
            v = v_[t, ...]
            if k in fr_dict.keys():
                fr_dict[k] += v.mean().item() / idx
            else:
                fr_dict[k] = v.mean().item() / idx
        return fr_dict

    model.eval()

    end = time.time()
    last_idx = len(loader) - 1
    fr_dict, nz_dict = {"t0": dict(), "t1": dict(), "t2": dict(), "t3": dict()}, {
        "t0": dict(),
        "t1": dict(),
        "t2": dict(),
        "t3": dict(),
    }
    with torch.no_grad():
        for batch_idx, (input, target) in enumerate(loader):
            last_batch = batch_idx == last_idx
            # if not args.prefetcher:
            input = input.cuda()
            target = target.cuda()
            if args.channels_last:
                input = input.contiguous(memory_format=torch.channels_last)

            with amp_autocast():
                output, firing_dict = model(input, hook=dict())
                if args.save_qkv and args.local_rank == 0:
                    torch.save(
                        firing_dict, os.path.join(output_dir, f"qkv_{batch_idx}.pkl")
                    )

            for t in range(args.time_steps):
                fr_single_dict = calc_firing_rate(
                    firing_dict, fr_dict["t" + str(t)], last_idx, t
                )
                fr_dict["t" + str(t)] = fr_single_dict
                nz_single_dict = calc_non_zero_rate(
                    firing_dict, nz_dict["t" + str(t)], last_idx, t
                )
                nz_dict["t" + str(t)] = nz_single_dict

            # augmentation reduction
            reduce_factor = args.tta
            if reduce_factor > 1:
                output = output.unfold(0, reduce_factor, reduce_factor).mean(dim=2)
                target = target[0 : target.size(0) : reduce_factor]

            loss = loss_fn(output, target)
            functional.reset_net(model)

            acc1, acc5 = accuracy(output, target, topk=(1, 5))

            if args.distributed:
                reduced_loss = reduce_tensor(loss.data, args.world_size)
                acc1 = reduce_tensor(acc1, args.world_size)
                acc5 = reduce_tensor(acc5, args.world_size)
            else:
                reduced_loss = loss.data

            torch.cuda.synchronize()

            losses_m.update(reduced_loss.item(), input.size(0))
            top1_m.update(acc1.item(), output.size(0))
            top5_m.update(acc5.item(), output.size(0))

            batch_time_m.update(time.time() - end)
            end = time.time()
            if args.local_rank == 0 and (
                last_batch or batch_idx % args.log_interval == 0
            ):
                log_name = "Test" + log_suffix
                _logger.info(
                    "{0}: [{1:>4d}/{2}]  "
                    "Time: {batch_time.val:.3f} ({batch_time.avg:.3f})  "
                    "Loss: {loss.val:>7.4f} ({loss.avg:>6.4f})  "
                    "Acc@1: {top1.val:>7.4f} ({top1.avg:>7.4f})  "
                    "Acc@5: {top5.val:>7.4f} ({top5.avg:>7.4f})".format(
                        log_name,
                        batch_idx,
                        last_idx,
                        batch_time=batch_time_m,
                        loss=losses_m,
                        top1=top1_m,
                        top5=top5_m,
                    )
                )

    metrics = OrderedDict(
        [
            ("loss", losses_m.avg),
            ("top1", top1_m.avg),
            ("top5", top5_m.avg),
            ("non_zero", nz_dict),
            ("firing_rate", fr_dict),
        ]
    )

    return metrics


if __name__ == "__main__":
    main()