Spaces:
Runtime error
Runtime error
File size: 6,513 Bytes
3b96cb1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 |
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import pickle
import shutil
import tempfile
import warnings
from typing import Optional, Sequence
import torch
from mmengine.dist import (barrier, broadcast, broadcast_object_list,
get_dist_info, is_main_process)
from mmengine.evaluator import BaseMetric
from mmengine.utils import mkdir_or_exist
class BaseVideoMetric(BaseMetric):
"""Base class for a metric in video task.
The metric first processes each batch of data_samples and predictions,
and appends the processed results to the results list. Then it
collects all results together from all ranks if distributed training
is used. Finally, it computes the metrics of the entire dataset.
A subclass of class:`BaseVideoMetric` should assign a meaningful value
to the class attribute `default_prefix`. See the argument `prefix` for
details.
"""
def process(self, data_batch: dict, data_samples: Sequence[dict]) -> None:
"""Process one batch of data samples and predictions.
The processed results should be stored in ``self.results``, which will
be used to compute the metrics when all batches have been processed.
Args:
data_batch (dict): A batch of data from the dataloader.
data_samples (Sequence[dict]): A batch of data samples that
contain annotations and predictions.
"""
for track_data_sample in data_samples:
video_data_samples = track_data_sample['video_data_samples']
ori_video_len = video_data_samples[0].ori_video_length
if ori_video_len == len(video_data_samples):
# video process
self.process_video(video_data_samples)
else:
# image process
self.process_image(video_data_samples, ori_video_len)
def evaluate(self, size: int = 1) -> dict:
"""Evaluate the model performance of the whole dataset after processing
all batches.
Args:
size (int): Length of the entire validation dataset.
Returns:
dict: Evaluation metrics dict on the val dataset. The keys are the
names of the metrics, and the values are corresponding results.
"""
if len(self.results) == 0:
warnings.warn(
f'{self.__class__.__name__} got empty `self.results`. Please '
'ensure that the processed results are properly added into '
'`self.results` in `process` method.')
results = collect_tracking_results(self.results, self.collect_device)
if is_main_process():
_metrics = self.compute_metrics(results) # type: ignore
# Add prefix to metric names
if self.prefix:
_metrics = {
'/'.join((self.prefix, k)): v
for k, v in _metrics.items()
}
metrics = [_metrics]
else:
metrics = [None] # type: ignore
broadcast_object_list(metrics)
# reset the results list
self.results.clear()
return metrics[0]
def collect_tracking_results(results: list,
device: str = 'cpu',
tmpdir: Optional[str] = None) -> Optional[list]:
"""Collected results in distributed environments. different from the
function mmengine.dist.collect_results, tracking compute metrics don't use
paramenter size, which means length of the entire validation dataset.
because it's equal to video num, but compute metrics need image num.
Args:
results (list): Result list containing result parts to be
collected. Each item of ``result_part`` should be a picklable
object.
device (str): Device name. Optional values are 'cpu' and 'gpu'.
tmpdir (str | None): Temporal directory for collected results to
store. If set to None, it will create a temporal directory for it.
``tmpdir`` should be None when device is 'gpu'. Defaults to None.
Returns:
list or None: The collected results.
"""
if device not in ['gpu', 'cpu']:
raise NotImplementedError(
f"device must be 'cpu' or 'gpu', but got {device}")
if device == 'gpu':
assert tmpdir is None, 'tmpdir should be None when device is "gpu"'
raise NotImplementedError('GPU collecting has not been supported yet')
else:
return collect_tracking_results_cpu(results, tmpdir)
def collect_tracking_results_cpu(result_part: list,
tmpdir: Optional[str] = None
) -> Optional[list]:
"""Collect results on cpu mode.
Saves the results on different gpus to 'tmpdir' and collects them by the
rank 0 worker.
Args:
result_part (list): The part of prediction results.
tmpdir (str): Path of directory to save the temporary results from
different gpus under cpu mode. If is None, use `tempfile.mkdtemp()`
to make a temporary path. Defaults to None.
Returns:
list or None: The collected results.
"""
rank, world_size = get_dist_info()
if world_size == 1:
return result_part
# create a tmp dir if it is not specified
if tmpdir is None:
MAX_LEN = 512
# 32 is whitespace
dir_tensor = torch.full((MAX_LEN, ), 32, dtype=torch.uint8)
if rank == 0:
mkdir_or_exist('.dist_test')
tmpdir = tempfile.mkdtemp(dir='.dist_test')
tmpdir = torch.tensor(
bytearray(tmpdir.encode()), dtype=torch.uint8)
dir_tensor[:len(tmpdir)] = tmpdir
broadcast(dir_tensor, 0)
tmpdir = dir_tensor.numpy().tobytes().decode().rstrip()
else:
mkdir_or_exist(tmpdir)
# dump the part result to the dir
with open(osp.join(tmpdir, f'part_{rank}.pkl'), 'wb') as f: # type: ignore
pickle.dump(result_part, f, protocol=2)
barrier()
# collect all parts
if rank != 0:
return None
else:
# load results of all parts from tmp dir
part_list = []
for i in range(world_size):
path = osp.join(tmpdir, f'part_{i}.pkl') # type: ignore
with open(path, 'rb') as f:
part_list.extend(pickle.load(f))
shutil.rmtree(tmpdir)
return part_list
|