# Copyright (c) OpenMMLab. All rights reserved. import math import os import os.path as osp import sys import warnings from bisect import bisect_right import mmcv import torch from mmcv.runner import HOOKS, Hook, get_dist_info from ..registry import build_metric @HOOKS.register_module() class GenerativeEvalHook(Hook): """Evaluation Hook for Generative Models. This evaluation hook can be used to evaluate unconditional and conditional models. Note that only ``FID`` and ``IS`` metric are supported for the distributed training now. In the future, we will support more metrics for the evaluation during the training procedure. In our config system, you only need to add `evaluation` with the detailed configureations. Below is several usage cases for different situations. What you need to do is to add these lines at the end of your config file. Then, you can use this evaluation hook in the training procedure. To be noted that, this evaluation hook support evaluation with dynamic intervals for FID or other metrics may fluctuate frequently at the end of the training process. # TODO: fix the online doc #. Only use FID for evaluation .. code-block:: python :linenos: evaluation = dict( type='GenerativeEvalHook', interval=10000, metrics=dict( type='FID', num_images=50000, inception_pkl='work_dirs/inception_pkl/ffhq-256-50k-rgb.pkl', bgr2rgb=True), sample_kwargs=dict(sample_model='ema')) #. Use FID and IS simultaneously and save the best checkpoints respectively .. code-block:: python :linenos: evaluation = dict( type='GenerativeEvalHook', interval=10000, metrics=[dict( type='FID', num_images=50000, inception_pkl='work_dirs/inception_pkl/ffhq-256-50k-rgb.pkl', bgr2rgb=True), dict(type='IS', num_images=50000)], best_metric=['fid', 'is'], sample_kwargs=dict(sample_model='ema')) #. Use dynamic evaluation intervals .. code-block:: python :linenos: # interval = 10000 if iter < 50000, # interval = 4000, if 50000 <= iter < 750000, # interval = 2000, if iter >= 750000 evaluation = dict( type='GenerativeEvalHook', interval=dict(milestones=[500000, 750000], interval=[10000, 4000, 2000]) metrics=[dict( type='FID', num_images=50000, inception_pkl='work_dirs/inception_pkl/ffhq-256-50k-rgb.pkl', bgr2rgb=True), dict(type='IS', num_images=50000)], best_metric=['fid', 'is'], sample_kwargs=dict(sample_model='ema')) Args: dataloader (DataLoader): A PyTorch dataloader. interval (int | dict): Evaluation interval. If int is passed, ``eval_hook`` would run under given interval. If a dict is passed, The key and value would be interpret as 'milestones' and 'interval' of the evaluation. Default: 1. dist (bool, optional): Whether to use distributed evaluation. Defaults to True. metrics (dict | list[dict], optional): Configs for metrics that will be used in evaluation hook. Defaults to None. sample_kwargs (dict | None, optional): Additional keyword arguments for sampling images. Defaults to None. save_best_ckpt (bool, optional): Whether to save the best checkpoint according to ``best_metric``. Defaults to ``True``. best_metric (str | list, optional): Which metric to be used in saving the best checkpoint. Multiple metrics have been supported by inputing a list of metric names, e.g., ``['fid', 'is']``. Defaults to ``'fid'``. """ rule_map = {'greater': lambda x, y: x > y, 'less': lambda x, y: x < y} init_value_map = {'greater': -math.inf, 'less': math.inf} greater_keys = ['acc', 'top', 'AR@', 'auc', 'precision', 'mAP', 'is'] less_keys = ['loss', 'fid'] _supported_best_metrics = ['fid', 'is'] def __init__(self, dataloader, interval=1, dist=True, metrics=None, sample_kwargs=None, save_best_ckpt=True, best_metric='fid'): assert metrics is not None self.dataloader = dataloader self.dist = dist self.sample_kwargs = sample_kwargs if sample_kwargs else dict() self.save_best_ckpt = save_best_ckpt self.best_metric = best_metric if isinstance(interval, int): self.interval = interval elif isinstance(interval, dict): if 'milestones' not in interval or 'interval' not in interval: raise KeyError( '`milestones` and `interval` must exist in interval dict ' 'if you want to use the dynamic interval evaluation ' f'strategy. But receive [{[k for k in interval.keys()]}] ' 'in the interval dict.') self.milestones = interval['milestones'] self.interval = interval['interval'] # check if length of interval match with the milestones if len(self.interval) != len(self.milestones) + 1: raise ValueError( f'Length of `interval`(={len(self.interval)}) cannot ' f'match length of `milestones`(={len(self.milestones)}).') # check if milestones is in order for idx in range(len(self.milestones) - 1): former, latter = self.milestones[idx], self.milestones[idx + 1] if former >= latter: raise ValueError( 'Elements in `milestones` should in ascending order.') else: raise TypeError('`interval` only support `int` or `dict`,' f'recieve {type(self.interval)} instead.') if isinstance(best_metric, str): self.best_metric = [self.best_metric] if self.save_best_ckpt: not_supported = set(self.best_metric) - set( self._supported_best_metrics) assert len(not_supported) == 0, ( f'{not_supported} is not supported for saving best ckpt') self.metrics = build_metric(metrics) if isinstance(metrics, dict): self.metrics = [self.metrics] for metric in self.metrics: metric.prepare() # add support for saving best ckpt if self.save_best_ckpt: self.rule = {} self.compare_func = {} self._curr_best_score = {} self._curr_best_ckpt_path = {} for name in self.best_metric: if name in self.greater_keys: self.rule[name] = 'greater' else: self.rule[name] = 'less' self.compare_func[name] = self.rule_map[self.rule[name]] self._curr_best_score[name] = self.init_value_map[ self.rule[name]] self._curr_best_ckpt_path[name] = None def get_current_interval(self, runner): """Get current evaluation interval. Args: runner (``mmcv.runner.BaseRunner``): The runner. """ if isinstance(self.interval, int): return self.interval else: curr_iter = runner.iter + 1 index = bisect_right(self.milestones, curr_iter) return self.interval[index] def before_run(self, runner): """The behavior before running. Args: runner (``mmcv.runner.BaseRunner``): The runner. """ if self.save_best_ckpt is not None: if runner.meta is None: warnings.warn('runner.meta is None. Creating an empty one.') runner.meta = dict() runner.meta.setdefault('hook_msgs', dict()) def after_train_iter(self, runner): """The behavior after each train iteration. Args: runner (``mmcv.runner.BaseRunner``): The runner. """ interval = self.get_current_interval(runner) if not self.every_n_iters(runner, interval): return runner.model.eval() batch_size = self.dataloader.batch_size rank, ws = get_dist_info() total_batch_size = batch_size * ws # sample real images max_real_num_images = max(metric.num_images - metric.num_real_feeded for metric in self.metrics) # define mmcv progress bar if rank == 0 and max_real_num_images > 0: mmcv.print_log( f'Sample {max_real_num_images} real images for evaluation', 'mmgen') pbar = mmcv.ProgressBar(max_real_num_images) if max_real_num_images > 0: for data in self.dataloader: if 'real_img' in data: reals = data['real_img'] # key for conditional GAN elif 'img' in data: reals = data['img'] else: raise KeyError('Cannot found key for images in data_dict. ' 'Only support `real_img` for unconditional ' 'datasets and `img` for conditional ' 'datasets.') if reals.shape[1] not in [1, 3]: raise RuntimeError('real images should have one or three ' 'channels in the first, ' 'not % d' % reals.shape[1]) if reals.shape[1] == 1: reals = reals.repeat(1, 3, 1, 1) num_feed = 0 for metric in self.metrics: num_feed_ = metric.feed(reals, 'reals') num_feed = max(num_feed_, num_feed) if num_feed <= 0: break if rank == 0: pbar.update(num_feed) max_num_images = max(metric.num_images for metric in self.metrics) if rank == 0: mmcv.print_log( f'Sample {max_num_images} fake images for evaluation', 'mmgen') # define mmcv progress bar if rank == 0: pbar = mmcv.ProgressBar(max_num_images) # sampling fake images and directly send them to metrics for _ in range(0, max_num_images, total_batch_size): with torch.no_grad(): fakes = runner.model( None, num_batches=batch_size, return_loss=False, **self.sample_kwargs) for metric in self.metrics: # feed in fake images metric.feed(fakes, 'fakes') if rank == 0: pbar.update(total_batch_size) runner.log_buffer.clear() # a dirty walkround to change the line at the end of pbar if rank == 0: sys.stdout.write('\n') for metric in self.metrics: with torch.no_grad(): metric.summary() for name, val in metric._result_dict.items(): runner.log_buffer.output[name] = val # record best metric and save the best ckpt if self.save_best_ckpt and name in self.best_metric: self._save_best_ckpt(runner, val, name) runner.log_buffer.ready = True runner.model.train() # clear all current states for next evaluation for metric in self.metrics: metric.clear() def _save_best_ckpt(self, runner, new_score, metric_name): """Save checkpoint with best metric score. Args: runner (``mmcv.runner.BaseRunner``): The runner. new_score (float): New metric score. metric_name (str): Name of metric. """ curr_iter = f'iter_{runner.iter + 1}' if self.compare_func[metric_name](new_score, self._curr_best_score[metric_name]): best_ckpt_name = f'best_{metric_name}_{curr_iter}.pth' runner.meta['hook_msgs'][f'best_score_{metric_name}'] = new_score if self._curr_best_ckpt_path[metric_name] and osp.isfile( self._curr_best_ckpt_path[metric_name]): os.remove(self._curr_best_ckpt_path[metric_name]) self._curr_best_ckpt_path[metric_name] = osp.join( runner.work_dir, best_ckpt_name) runner.save_checkpoint( runner.work_dir, best_ckpt_name, create_symlink=False) runner.meta['hook_msgs'][ f'best_ckpt_{metric_name}'] = self._curr_best_ckpt_path[ metric_name] self._curr_best_score[metric_name] = new_score runner.logger.info( f'Now best checkpoint is saved as {best_ckpt_name}.') runner.logger.info(f'Best {metric_name} is {new_score:0.4f} ' f'at {curr_iter}.') @HOOKS.register_module() class TranslationEvalHook(GenerativeEvalHook): """Evaluation Hook for Translation Models. This evaluation hook can be used to evaluate translation models. Note that only ``FID`` and ``IS`` metric are supported for the distributed training now. In the future, we will support more metrics for the evaluation during the training procedure. In our config system, you only need to add `evaluation` with the detailed configureations. Below is several usage cases for different situations. What you need to do is to add these lines at the end of your config file. Then, you can use this evaluation hook in the training procedure. To be noted that, this evaluation hook support evaluation with dynamic intervals for FID or other metrics may fluctuate frequently at the end of the training process. # TODO: fix the online doc #. Only use FID for evaluation .. code-blcok:: python :linenos evaluation = dict( type='TranslationEvalHook', target_domain='photo', interval=10000, metrics=dict(type='FID', num_images=106, bgr2rgb=True)) #. Use FID and IS simultaneously and save the best checkpoints respectively .. code-block:: python :linenos evaluation = dict( type='TranslationEvalHook', target_domain='photo', interval=10000, metrics=[ dict(type='FID', num_images=106, bgr2rgb=True), dict( type='IS', num_images=106, inception_args=dict(type='pytorch')) ], best_metric=['fid', 'is']) #. Use dynamic evaluation intervals .. code-block:: python :linenos # interval = 10000 if iter < 100000, # interval = 4000, if 100000 <= iter < 200000, # interval = 2000, if iter >= 200000 evaluation = dict( type='TranslationEvalHook', interval=dict(milestones=[100000, 200000], interval=[10000, 4000, 2000]), target_domain='zebra', metrics=[ dict(type='FID', num_images=140, bgr2rgb=True), dict(type='IS', num_images=140) ], best_metric=['fid', 'is']) Args: target_domain (str): Target domain of output image. """ def __init__(self, *args, target_domain, **kwargs): super().__init__(*args, **kwargs) self.target_domain = target_domain def after_train_iter(self, runner): """The behavior after each train iteration. Args: runner (``mmcv.runner.BaseRunner``): The runner. """ interval = self.get_current_interval(runner) if not self.every_n_iters(runner, interval): return runner.model.eval() source_domain = runner.model.module.get_other_domains( self.target_domain)[0] # feed real images max_num_images = max(metric.num_images for metric in self.metrics) for metric in self.metrics: if metric.num_real_feeded >= metric.num_real_need: continue mmcv.print_log(f'Feed reals to {metric.name} metric.', 'mmgen') # feed in real images for data in self.dataloader: # key for translation model if f'img_{self.target_domain}' in data: reals = data[f'img_{self.target_domain}'] # key for conditional GAN else: raise KeyError( 'Cannot found key for images in data_dict. ') num_feed = metric.feed(reals, 'reals') if num_feed <= 0: break mmcv.print_log(f'Sample {max_num_images} fake images for evaluation', 'mmgen') rank, ws = get_dist_info() # define mmcv progress bar if rank == 0: pbar = mmcv.ProgressBar(max_num_images) # feed in fake images for data in self.dataloader: # key for translation model if f'img_{source_domain}' in data: with torch.no_grad(): output_dict = runner.model( data[f'img_{source_domain}'], test_mode=True, target_domain=self.target_domain, **self.sample_kwargs) fakes = output_dict['target'] # key Error else: raise KeyError('Cannot found key for images in data_dict. ') # sampling fake images and directly send them to metrics # pbar update number for one proc num_update = 0 for metric in self.metrics: if metric.num_fake_feeded >= metric.num_fake_need: continue num_feed = metric.feed(fakes, 'fakes') num_update = max(num_update, num_feed) if num_feed <= 0: break if rank == 0: if num_update > 0: pbar.update(num_update * ws) runner.log_buffer.clear() # a dirty walkround to change the line at the end of pbar if rank == 0: sys.stdout.write('\n') for metric in self.metrics: with torch.no_grad(): metric.summary() for name, val in metric._result_dict.items(): runner.log_buffer.output[name] = val # record best metric and save the best ckpt if self.save_best_ckpt and name in self.best_metric: self._save_best_ckpt(runner, val, name) runner.log_buffer.ready = True runner.model.train() # clear all current states for next evaluation for metric in self.metrics: metric.clear()