Spaces:
Running
Running
| #!/usr/bin/env python3 -u | |
| # Copyright 2022 The OFA-Sys Team. | |
| # All rights reserved. | |
| # This source code is licensed under the Apache 2.0 license | |
| # found in the LICENSE file in the root directory. | |
| import logging | |
| import os | |
| import sys | |
| import numpy as np | |
| import torch | |
| from fairseq import distributed_utils, options, tasks, utils | |
| from fairseq.dataclass.utils import convert_namespace_to_omegaconf | |
| from fairseq.logging import progress_bar | |
| from fairseq.utils import reset_logging | |
| from omegaconf import DictConfig | |
| from utils import checkpoint_utils | |
| from utils.eval_utils import eval_step, merge_results | |
| logging.basicConfig( | |
| format="%(asctime)s | %(levelname)s | %(name)s | %(message)s", | |
| datefmt="%Y-%m-%d %H:%M:%S", | |
| level=os.environ.get("LOGLEVEL", "INFO").upper(), | |
| stream=sys.stdout, | |
| ) | |
| logger = logging.getLogger("ofa.evaluate") | |
| def apply_half(t): | |
| if t.dtype is torch.float32: | |
| return t.to(dtype=torch.half) | |
| return t | |
| def main(cfg: DictConfig, **kwargs): | |
| utils.import_user_module(cfg.common) | |
| reset_logging() | |
| logger.info(cfg) | |
| assert ( | |
| cfg.dataset.max_tokens is not None or cfg.dataset.batch_size is not None | |
| ), "Must specify batch size either with --max-tokens or --batch-size" | |
| # Fix seed for stochastic decoding | |
| if cfg.common.seed is not None and not cfg.generation.no_seed_provided: | |
| np.random.seed(cfg.common.seed) | |
| utils.set_torch_seed(cfg.common.seed) | |
| use_fp16 = cfg.common.fp16 | |
| use_cuda = torch.cuda.is_available() and not cfg.common.cpu | |
| if use_cuda: | |
| torch.cuda.set_device(cfg.distributed_training.device_id) | |
| # Load ensemble | |
| overrides = eval(cfg.common_eval.model_overrides) | |
| # Deal with beam-search / all-candidate VQA eval | |
| if cfg.task._name == "vqa_gen": | |
| overrides['val_inference_type'] = "beamsearch" if kwargs['beam_search_vqa_eval'] else "allcand" | |
| logger.info("loading model(s) from {}".format(cfg.common_eval.path)) | |
| if kwargs["zero_shot"]: | |
| task = tasks.setup_task(cfg.task) | |
| models, saved_cfg = checkpoint_utils.load_model_ensemble( | |
| utils.split_paths(cfg.common_eval.path), | |
| arg_overrides=overrides, | |
| task=task, | |
| suffix=cfg.checkpoint.checkpoint_suffix, | |
| strict=(cfg.checkpoint.checkpoint_shard_count == 1), | |
| num_shards=cfg.checkpoint.checkpoint_shard_count, | |
| ) | |
| else: | |
| models, saved_cfg, task = checkpoint_utils.load_model_ensemble_and_task( | |
| utils.split_paths(cfg.common_eval.path), | |
| arg_overrides=overrides, | |
| suffix=cfg.checkpoint.checkpoint_suffix, | |
| strict=(cfg.checkpoint.checkpoint_shard_count == 1), | |
| num_shards=cfg.checkpoint.checkpoint_shard_count, | |
| ) | |
| # loading the dataset should happen after the checkpoint has been loaded so we can give it the saved task config | |
| task.load_dataset(cfg.dataset.gen_subset, task_cfg=saved_cfg.task) | |
| # Move models to GPU | |
| for model, ckpt_path in zip(models, utils.split_paths(cfg.common_eval.path)): | |
| if kwargs['ema_eval']: | |
| logger.info("loading EMA weights from {}".format(ckpt_path)) | |
| model.load_state_dict(checkpoint_utils.load_ema_from_checkpoint(ckpt_path)['model']) | |
| model.eval() | |
| if use_fp16: | |
| model.half() | |
| if use_cuda and not cfg.distributed_training.pipeline_model_parallel: | |
| model.cuda() | |
| model.prepare_for_inference_(cfg) | |
| # Load dataset (possibly sharded) | |
| itr = task.get_batch_iterator( | |
| dataset=task.dataset(cfg.dataset.gen_subset), | |
| max_tokens=cfg.dataset.max_tokens, | |
| max_sentences=cfg.dataset.batch_size, | |
| max_positions=utils.resolve_max_positions( | |
| task.max_positions(), *[m.max_positions() for m in models] | |
| ), | |
| ignore_invalid_inputs=cfg.dataset.skip_invalid_size_inputs_valid_test, | |
| required_batch_size_multiple=cfg.dataset.required_batch_size_multiple, | |
| seed=cfg.common.seed, | |
| num_shards=cfg.distributed_training.distributed_world_size, | |
| shard_id=cfg.distributed_training.distributed_rank, | |
| num_workers=cfg.dataset.num_workers, | |
| data_buffer_size=cfg.dataset.data_buffer_size, | |
| ).next_epoch_itr(shuffle=False) | |
| progress = progress_bar.progress_bar( | |
| itr, | |
| log_format=cfg.common.log_format, | |
| log_interval=cfg.common.log_interval, | |
| default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"), | |
| ) | |
| # Initialize generator | |
| generator = task.build_generator(models, cfg.generation) | |
| # for sample in progress: | |
| # if "net_input" not in sample: | |
| # continue | |
| # sample = utils.move_to_cuda(sample) if use_cuda else sample | |
| # sample = utils.apply_to_sample(apply_half, sample) if cfg.common.fp16 else sample | |
| # with torch.no_grad(): | |
| # eval_step(task, generator, models, sample, **kwargs) | |
| # progress.log({"sentences": sample["nsentences"]}) | |
| # | |
| # merge_results(task, cfg, logger, kwargs['result_dir']) | |
| results = [] | |
| prec_list = [.5, .6, .7, .8, .9] | |
| prec_score_sum = [torch.FloatTensor([0]).cuda() for _ in prec_list] | |
| f_score_sum = torch.FloatTensor([0]).cuda() | |
| ap_det_score_sum = torch.FloatTensor([0]).cuda() | |
| score_sum = torch.FloatTensor([0]).cuda() | |
| score_cnt = torch.FloatTensor([0]).cuda() | |
| cum_I_sum = torch.FloatTensor([0]).cuda() | |
| cum_U_sum = torch.FloatTensor([0]).cuda() | |
| for sample in progress: | |
| if "net_input" not in sample: | |
| continue | |
| sample = utils.move_to_cuda(sample) if use_cuda else sample | |
| sample = utils.apply_to_sample(apply_half, sample) if cfg.common.fp16 else sample | |
| with torch.no_grad(): | |
| result, scores, f_scores, ap_scores, cum_I, cum_U = eval_step(task, generator, models, sample, **kwargs) | |
| results += result | |
| for prec_score, prec in zip(prec_score_sum, prec_list): | |
| prec_score += sum(scores >= prec) if scores is not None else 0 | |
| cum_I_sum += sum(cum_I) if scores is not None else 0 | |
| cum_U_sum += sum(cum_U) if scores is not None else 0 | |
| score_sum += sum(scores) if scores is not None else 0 | |
| f_score_sum += sum(f_scores) if scores is not None else 0 | |
| ap_det_score_sum += sum(ap_scores) if scores is not None else 0 | |
| score_cnt += len(scores) if scores is not None else 0 | |
| progress.log({"sentences": sample["nsentences"]}) | |
| merge_results(task, cfg, logger, score_cnt, score_sum, f_score_sum, ap_det_score_sum,prec_score_sum, cum_I_sum, cum_U_sum, results) | |
| def cli_main(): | |
| parser = options.get_generation_parser() | |
| parser.add_argument("--ema-eval", action='store_true', help="Use EMA weights to make evaluation.") | |
| parser.add_argument("--beam-search-vqa-eval", action='store_true', help="Use beam search for vqa evaluation (faster inference speed but sub-optimal result), if not specified, we compute scores for each answer in the candidate set, which is slower but can obtain best result.") | |
| parser.add_argument("--zero-shot", action='store_true') | |
| parser.add_argument("--vis_dir", type=str, default=None) | |
| parser.add_argument("--result_dir", type=str, default=None) | |
| parser.add_argument("--vis", action='store_true', default=False) | |
| args = options.parse_args_and_arch(parser) | |
| cfg = convert_namespace_to_omegaconf(args) | |
| if args.result_dir is None: | |
| args.result_dir = args.vis_dir | |
| distributed_utils.call_main( | |
| cfg, main, ema_eval=args.ema_eval, beam_search_vqa_eval=args.beam_search_vqa_eval, zero_shot=args.zero_shot, | |
| vis_dir=args.vis_dir, vis=args.vis, result_dir=args.result_dir | |
| ) | |
| if __name__ == "__main__": | |
| cli_main() | |