|
|
|
|
| import torch
|
| from mmengine.model import is_model_wrapper
|
| from mmengine.runner import TestLoop, ValLoop, autocast
|
|
|
| from mmaction.registry import LOOPS
|
|
|
|
|
| @LOOPS.register_module()
|
| class RetrievalValLoop(ValLoop):
|
| """Loop for multimodal retrieval val.
|
|
|
| Args:
|
| runner (Runner): A reference of runner.
|
| dataloader (Dataloader or dict): A dataloader object or a dict to
|
| build a dataloader.
|
| evaluator (Evaluator or dict or list): Used for computing metrics.
|
| fp16 (bool): Whether to enable fp16 valing. Defaults to
|
| False.
|
| """
|
|
|
| def run(self) -> dict:
|
| """Launch val."""
|
| self.runner.call_hook('before_val')
|
| self.runner.call_hook('before_val_epoch')
|
| self.runner.model.eval()
|
|
|
| feats_local = []
|
| data_samples_local = []
|
|
|
| for idx, data_batch in enumerate(self.dataloader):
|
| with torch.no_grad():
|
| self.runner.call_hook(
|
| 'before_val_iter', batch_idx=idx, data_batch=data_batch)
|
|
|
| with autocast(enabled=self.fp16):
|
| if is_model_wrapper(self.runner.model):
|
| data_preprocessor = self.runner.model.module.data_preprocessor
|
| else:
|
| data_preprocessor = self.runner.model.data_preprocessor
|
|
|
|
|
| data_batch = data_preprocessor(data_batch, False)
|
| feats = self.runner.model._run_forward(
|
| data_batch, mode='tensor')
|
| feats_local.append(feats)
|
| data_samples_local.extend(data_batch['data_samples'])
|
| self.runner.call_hook(
|
| 'after_val_iter',
|
| batch_idx=idx,
|
| data_batch=data_batch,
|
| outputs=feats)
|
|
|
|
|
| feats_local = {
|
| k: torch.cat([dic[k] for dic in feats_local])
|
| for k in feats_local[0]
|
| }
|
|
|
|
|
| if is_model_wrapper(self.runner.model):
|
| predict_all_fn = self.runner.model.module.predict_all
|
| else:
|
| predict_all_fn = self.runner.model.predict_all
|
|
|
| num_videos = self.dataloader.dataset.num_videos
|
| num_texts = self.dataloader.dataset.num_texts
|
| with torch.no_grad():
|
| with autocast(enabled=self.fp16):
|
| i2t_data_samples, t2i_data_samples = predict_all_fn(
|
| feats_local,
|
| data_samples_local,
|
| num_images=num_videos,
|
| num_texts=num_texts,
|
| )
|
|
|
| self.evaluator.process(i2t_data_samples, None)
|
| i2t_metrics = self.evaluator.evaluate(num_videos)
|
| i2t_metrics = {f'i2t/{k}': v for k, v in i2t_metrics.items()}
|
| self.evaluator.process(t2i_data_samples, None)
|
| t2i_metrics = self.evaluator.evaluate(num_texts)
|
| t2i_metrics = {f't2i/{k}': v for k, v in t2i_metrics.items()}
|
| metrics = {**i2t_metrics, **t2i_metrics}
|
| self.runner.call_hook('after_val_epoch', metrics=metrics)
|
| self.runner.call_hook('after_val')
|
| return metrics
|
|
|
|
|
| @LOOPS.register_module()
|
| class RetrievalTestLoop(TestLoop):
|
| """Loop for multimodal retrieval test.
|
|
|
| Args:
|
| runner (Runner): A reference of runner.
|
| dataloader (Dataloader or dict): A dataloader object or a dict to
|
| build a dataloader.
|
| evaluator (Evaluator or dict or list): Used for computing metrics.
|
| fp16 (bool): Whether to enable fp16 testing. Defaults to
|
| False.
|
| """
|
|
|
| def run(self) -> dict:
|
| """Launch test."""
|
| self.runner.call_hook('before_test')
|
| self.runner.call_hook('before_test_epoch')
|
| self.runner.model.eval()
|
|
|
| feats_local = []
|
| data_samples_local = []
|
|
|
| for idx, data_batch in enumerate(self.dataloader):
|
| with torch.no_grad():
|
| self.runner.call_hook(
|
| 'before_test_iter', batch_idx=idx, data_batch=data_batch)
|
|
|
| with autocast(enabled=self.fp16):
|
| if is_model_wrapper(self.runner.model):
|
| data_preprocessor = self.runner.model.module.data_preprocessor
|
| else:
|
| data_preprocessor = self.runner.model.data_preprocessor
|
|
|
| data_batch = data_preprocessor(data_batch, False)
|
| feats = self.runner.model._run_forward(
|
| data_batch, mode='tensor')
|
| feats_local.append(feats)
|
| data_samples_local.extend(data_batch['data_samples'])
|
| self.runner.call_hook(
|
| 'after_test_iter',
|
| batch_idx=idx,
|
| data_batch=data_batch,
|
| outputs=feats)
|
|
|
|
|
| feats_local = {
|
| k: torch.cat([dic[k] for dic in feats_local])
|
| for k in feats_local[0]
|
| }
|
|
|
|
|
| if is_model_wrapper(self.runner.model):
|
| predict_all_fn = self.runner.model.module.predict_all
|
| else:
|
| predict_all_fn = self.runner.model.predict_all
|
|
|
| num_videos = self.dataloader.dataset.num_videos
|
| num_texts = self.dataloader.dataset.num_texts
|
| with torch.no_grad():
|
| with autocast(enabled=self.fp16):
|
| i2t_data_samples, t2i_data_samples = predict_all_fn(
|
| feats_local,
|
| data_samples_local,
|
| num_images=num_videos,
|
| num_texts=num_texts,
|
| )
|
|
|
|
|
| self.evaluator.process(i2t_data_samples, None)
|
| i2t_metrics = self.evaluator.evaluate(num_videos)
|
| i2t_metrics = {f'i2t/{k}': v for k, v in i2t_metrics.items()}
|
| self.evaluator.process(t2i_data_samples, None)
|
| t2i_metrics = self.evaluator.evaluate(num_texts)
|
| t2i_metrics = {f't2i/{k}': v for k, v in t2i_metrics.items()}
|
| metrics = {**i2t_metrics, **t2i_metrics}
|
|
|
| self.runner.call_hook('after_test_epoch', metrics=metrics)
|
| self.runner.call_hook('after_test')
|
| return metrics
|
|
|
