finetune_validate.py

# Copyright (c) 2022 Graphcore Ltd. All rights reserved.
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# References:
# DeiT: https://github.com/facebookresearch/deit
# BEiT: https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------

import time
import numpy as np
import os
from pathlib import Path
import torch
import poptorch
import yaml
from util.ipu_mixup import Mixup
from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy
from util.datasets import build_dataset
import core.models_vit as models_vit
from options import finetune_options
from core import utils
from util.log import AverageMeter, ProgressMeter, logger, WandbLog
from timm.utils import accuracy
import logging
from argparser import get_args_parser
import wandb


@torch.no_grad()
def evaluate(data_loader, model, print_freq):
    criterion = torch.nn.CrossEntropyLoss()

    batch_time = AverageMeter("time", ":6.3f")
    data_time = AverageMeter("data", ":6.3f")
    losses = AverageMeter("loss", ":.3f")
    tput = AverageMeter("throughput", ":.0f")
    acc1_log = AverageMeter("acc1", ":.3f")
    acc5_log = AverageMeter("acc5", ":.3f")
    # switch to evaluation mode
    model.eval()
    meters = [batch_time, data_time, losses, tput, acc1_log, acc5_log]
    progress = ProgressMeter(len(data_loader), meters, prefix="Evaluation: ")

    if args.wandb:
        wandb_logger = WandbLog(meters)

    end = time.time()
    for it, batch in enumerate(data_loader):
        data_time.update(time.time() - end)
        print(data_time.val)
        images = batch[0]
        target = batch[-1]
        # compute output
        output = model(images)
        loss = criterion(output, target)
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        batch_size = images.shape[0]

        acc1 = torch.mean(acc1)
        acc5 = torch.mean(acc5)

        acc1_log.update(acc1.item(), batch_size)
        acc5_log.update(acc5.item(), batch_size)
        losses.update(loss.item(), batch_size)
        batch_time.update(time.time() - end)

        end = time.time()

        tput.update(utils.sync_metrics(batch_size / batch_time.val))

        if it % print_freq == 0:
            if not args.use_popdist or (args.use_popdist and args.popdist_rank == 0):
                logger.info(progress.display(it))
                if args.wandb:
                    wandb_logger.log()

    logger.info("* Acc@1 {top1} Acc@5 {top5} loss {losses}".format(top1=acc1_log, top5=acc5_log, losses=losses))

    return {meter.name: meter.avg for meter in meters}


def load_eval_checkpoint(model, path):
    assert os.path.exists(path), f"{path} not exists"
    model_state = torch.load(path)
    weights = model_state["model"]

    model.load_state_dict(weights)
    logger.info(f"Loaded checkpoint from path: {path}")
    return model


def main(args):
    log_path = os.path.join(args.output, "eval")
    now = time.strftime("%Y-%m-%d-%H_%M_%S", time.localtime(time.time()))
    fileHandler = logging.FileHandler(log_path + "_" + now + ".log", mode="w", encoding="UTF-8")
    fileHandler.setLevel(logging.NOTSET)
    logger.addHandler(fileHandler)
    args.async_type = "async"
    args.replica = 1
    args.batch_size = 16
    args.di = 1000
    opts_infer = finetune_options(
        gradient_accumulation_count=args.gradient_accumulation_count,
        replica=args.replica,
        half=args.half,
        ipu_per_replica=args.ipus,
        device_iterations=args.device_iterations,
        opt_type="eval",
    )

    logger.info("job dir: {}".format(os.path.dirname(os.path.realpath(__file__))))
    logger.info("{}".format(args).replace(", ", ",\n"))
    # fix the seed for reproducibility
    seed = args.seed
    torch.manual_seed(seed)
    np.random.seed(seed)

    dataset_val = build_dataset(is_train=False, args=args)

    if args.async_type == "async":
        mode = poptorch.DataLoaderMode.Async
        data_loader_val = poptorch.DataLoader(
            options=opts_infer,
            dataset=dataset_val,
            batch_size=args.batch_size,
            num_workers=args.num_workers,
            shuffle=False,
            drop_last=True,
            persistent_workers=True,
            worker_init_fn=None,
            mode=mode,
            async_options={"load_indefinitely": True},
        )
    elif args.async_type == "rebatch":
        mode = poptorch.DataLoaderMode.AsyncRebatched

        data_loader_val = poptorch.DataLoader(
            options=opts_infer,
            dataset=dataset_val,
            batch_size=args.batch_size,
            num_workers=args.num_workers,
            shuffle=False,
            drop_last=True,
            persistent_workers=True,
            worker_init_fn=None,
            mode=mode,
            async_options={"load_indefinitely": True},
            rebatched_worker_size=args.rebatched_worker_size,
        )
    else:
        data_loader_val = poptorch.DataLoader(
            options=opts_infer,
            dataset=dataset_val,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=args.num_workers,
            drop_last=True,
        )

    mixup_fn = None
    mixup_active = args.mixup > 0 or args.cutmix > 0.0 or args.cutmix_minmax is not None
    if mixup_active:
        logger.info("Mixup is activated!")
        mixup_fn = Mixup(
            mixup_alpha=args.mixup,
            cutmix_alpha=args.cutmix,
            cutmix_minmax=args.cutmix_minmax,
            prob=args.mixup_prob,
            switch_prob=args.mixup_switch_prob,
            mode=args.mixup_mode,
            label_smoothing=args.smoothing,
            num_classes=args.nb_classes,
        )

    if mixup_fn is not None:
        # smoothing is handled with mixup label transform
        criterion = SoftTargetCrossEntropy()
    elif args.smoothing > 0.0:
        criterion = LabelSmoothingCrossEntropy(smoothing=args.smoothing)
    else:
        criterion = torch.nn.CrossEntropyLoss()

    model = models_vit.__dict__[args.model](
        criterion=criterion,
        pipeline=[3, 3, 3, 3],
        num_classes=args.nb_classes,
        drop_path_rate=args.drop_path,
        global_pool=args.global_pool,
    )

    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
    logger.info("number of params (M): %.2f" % (n_parameters / 1.0e6))

    model = load_eval_checkpoint(model, args.resume)

    # switch to evaluation mode before creating an inference model
    model.eval()
    ipu_infer_model = poptorch.inferenceModel(model, options=opts_infer)
    test_stats = evaluate(data_loader_val, ipu_infer_model, args.print_freq)
    logger.info(f"Accuracy of the network on the {len(data_loader_val)} test images: {test_stats['acc1']:.1f}%")


if __name__ == "__main__":
    args = get_args_parser()

    utils.init_popdist(args)

    if args.wandb:
        if not args.use_popdist or (args.use_popdist and args.popdist_rank == 0):
            wandb.init(
                project=args.wandb_project_name,
                name=args.wandb_run_name,
                settings=wandb.Settings(console="wrap"),
                config=vars(args),
            )
            if args.half:
                wandb.config.update({"precision": "16.16"})
            else:
                wandb.config.update({"precision": "32.32"})

    if args.output:
        Path(args.output).mkdir(parents=True, exist_ok=True)
    main(args)