import os import sys from pathlib import Path if (_package_root := str(Path(__file__).absolute().parents[2])) not in sys.path: sys.path.insert(0, _package_root) import json from typing import * import importlib import importlib.util import click def _load_baseline(baseline_code_path: str, extra_args: Sequence[str]) -> 'MGEBaselineInterface': """Load the baseline model from its python file, forwarding the unparsed CLI args to its `load` command.""" from moge.utils.tools import import_file_as_module module = import_file_as_module(baseline_code_path, Path(baseline_code_path).stem) baseline_cls: Type['MGEBaselineInterface'] = getattr(module, 'Baseline') return baseline_cls.load.main(list(extra_args), standalone_mode=False) def _evaluate_benchmarks( baseline: 'MGEBaselineInterface', benchmarks: Iterable[Tuple[str, Dict[str, Any]]], *, metrics_output_path: Union[str, Path], dump_dir: Union[str, Path], oracle_mode: bool = False, mg: Optional[str] = None, dump_pred: bool = False, dump_gt: bool = False, tqdm_position: Optional[int] = None, tqdm_prefix: str = '', ) -> Dict[str, Any]: """Evaluate `baseline` on `benchmarks` sequentially, returning `{benchmark_name: metrics}` (without `mean`). `metrics_output_path` is where this process saves its own (possibly partial) results; `dump_dir` is the root for `--dump_pred` / `--dump_gt` outputs and always derives from the user-supplied `--output`. """ # Lazy import import cv2 import numpy as np from tqdm import tqdm import torch from moge.test.dataloader import EvalDataLoaderPipeline from moge.test.metrics import compute_metrics from moge.utils.geometry_torch import intrinsics_to_fov from moge.utils.vis import colorize_depth, colorize_normal from moge.utils.tools import key_average, timeit all_metrics = {} # A worker claims its benchmarks one at a time, so only the single-process path knows the total upfront. if tqdm_position is None: benchmarks = tqdm(list(benchmarks), desc='Benchmarks') # Iterate over the dataset for benchmark_name, benchmark_config in benchmarks: metrics_list = [] with ( EvalDataLoaderPipeline(**benchmark_config) as eval_data_pipe, tqdm(total=len(eval_data_pipe), desc=f'{tqdm_prefix}{benchmark_name}', position=tqdm_position, leave=False) as pbar ): # Iterate over the samples in the dataset for i in range(len(eval_data_pipe)): sample = eval_data_pipe.get() sample = {k: v.to(baseline.device) if isinstance(v, torch.Tensor) else v for k, v in sample.items()} image = sample['image'] gt_intrinsics = sample['intrinsics'] # Inference torch.cuda.synchronize() with torch.inference_mode(), timeit('_inference_timer', verbose=False) as timer: if oracle_mode: pred = baseline.infer_for_evaluation(image, gt_intrinsics) else: pred = baseline.infer_for_evaluation(image) torch.cuda.synchronize() # Compute metrics metrics, misc = compute_metrics(pred, sample, vis=dump_pred or dump_gt, mg=mg) metrics['inference_time'] = timer.time metrics_list.append(metrics) # Dump results dump_path = Path(dump_dir, f'{benchmark_name}', sample['filename'].replace('.zip', '')) if dump_pred: dump_path.joinpath('pred').mkdir(parents=True, exist_ok=True) cv2.imwrite(str(dump_path / 'pred' / 'image.jpg'), cv2.cvtColor((image.cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)) with Path(dump_path, 'pred', 'metrics.json').open('w') as f: json.dump(metrics, f, indent=4) if 'pred_points' in misc: points = misc['pred_points'].cpu().numpy() cv2.imwrite(str(dump_path / 'pred' / 'points.exr'), cv2.cvtColor(points.astype(np.float32), cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) if 'pred_depth' in misc: depth = misc['pred_depth'].cpu().numpy() if 'mask' in pred: mask = pred['mask'].cpu().numpy() depth = np.where(mask, depth, np.inf) cv2.imwrite(str(dump_path / 'pred' / 'depth.png'), cv2.cvtColor(colorize_depth(depth), cv2.COLOR_RGB2BGR)) if 'mask' in pred: mask = pred['mask'].cpu().numpy() cv2.imwrite(str(dump_path / 'pred' / 'mask.png'), (mask * 255).astype(np.uint8)) if 'normal' in pred: normal = pred['normal'].cpu().numpy() cv2.imwrite(str(dump_path / 'pred' / 'normal.png'), cv2.cvtColor(colorize_normal(normal), cv2.COLOR_RGB2BGR)) if 'intrinsics' in pred: intrinsics = pred['intrinsics'] fov_x, fov_y = intrinsics_to_fov(intrinsics) with open(dump_path / 'pred' / 'fov.json', 'w') as f: json.dump({ 'fov_x': np.rad2deg(fov_x.item()), 'fov_y': np.rad2deg(fov_y.item()), 'intrinsics': intrinsics.cpu().numpy().tolist(), }, f) if dump_gt: dump_path.joinpath('gt').mkdir(parents=True, exist_ok=True) cv2.imwrite(str(dump_path / 'gt' / 'image.jpg'), cv2.cvtColor((image.cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)) if 'points' in sample: points = sample['points'] cv2.imwrite(str(dump_path / 'gt' / 'points.exr'), cv2.cvtColor(points.cpu().numpy().astype(np.float32), cv2.COLOR_RGB2BGR), [cv2.IMWRITE_EXR_TYPE, cv2.IMWRITE_EXR_TYPE_FLOAT]) if 'depth' in sample: depth = sample['depth'] mask = sample['depth_mask'] cv2.imwrite(str(dump_path / 'gt' / 'depth.png'), cv2.cvtColor(colorize_depth(depth.cpu().numpy(), mask=mask.cpu().numpy()), cv2.COLOR_RGB2BGR)) if 'normal' in sample: normal = sample['normal'] cv2.imwrite(str(dump_path / 'gt' / 'normal.png'), cv2.cvtColor(colorize_normal(normal.cpu().numpy()), cv2.COLOR_RGB2BGR)) if 'depth_mask' in sample: mask = sample['depth_mask'] cv2.imwrite(str(dump_path / 'gt' /'mask.png'), (mask.cpu().numpy() * 255).astype(np.uint8)) if 'intrinsics' in sample: intrinsics = sample['intrinsics'] fov_x, fov_y = intrinsics_to_fov(intrinsics) with open(dump_path / 'gt' / 'info.json', 'w') as f: json.dump({ 'fov_x': np.rad2deg(fov_x.item()), 'fov_y': np.rad2deg(fov_y.item()), 'intrinsics': intrinsics.cpu().numpy().tolist(), }, f) # Save intermediate results if i % 100 == 0 or i == len(eval_data_pipe) - 1: Path(metrics_output_path).write_text( json.dumps({ **all_metrics, benchmark_name: key_average(metrics_list) }, indent=4) ) pbar.update(1) all_metrics[benchmark_name] = key_average(metrics_list) return all_metrics def _resolve_visible_devices(ngpu: int) -> List[str]: """Return `ngpu` device tokens to be assigned to the workers as their `CUDA_VISIBLE_DEVICES`.""" cuda_visible_devices = os.environ.get('CUDA_VISIBLE_DEVICES') if cuda_visible_devices is not None: # NOTE: entries of CUDA_VISIBLE_DEVICES are absolute ids (and may be GPU UUIDs or MIG ids), # so they are sliced as opaque strings and never renumbered. devices = [d.strip() for d in cuda_visible_devices.split(',') if d.strip()] source = 'CUDA_VISIBLE_DEVICES' else: import torch # queried via NVML, does not create a CUDA context devices = [str(i) for i in range(torch.cuda.device_count())] source = 'torch.cuda.device_count()' if not devices: raise click.UsageError(f'No CUDA device is available ({source} reports none).') if ngpu > len(devices): raise click.UsageError(f'--ngpu {ngpu} exceeds the {len(devices)} available GPU(s) ({source}: {",".join(devices)}).') return devices[:ngpu] def _partial_output_path(output_path: Union[str, Path], rank: int) -> Path: """`eval_output/moge.json` -> `eval_output/moge.rank0.json`""" path = Path(output_path) return path.with_name(f'{path.stem}.rank{rank}{path.suffix or ".json"}') def _claim_benchmarks(benchmarks: Sequence[Tuple[str, Dict[str, Any]]], counter) -> Iterator[Tuple[str, Dict[str, Any]]]: """Yield benchmarks claimed one at a time from a counter shared by all workers. Whoever finishes first takes the next benchmark, so the workers stay busy no matter how unevenly the benchmarks are sized. A plain shared counter is used rather than a queue because a queue's feeder thread may not have flushed yet when a worker first polls it, which would make that worker give up before any work is visible. """ while True: with counter.get_lock(): index = counter.value counter.value += 1 if index >= len(benchmarks): return yield benchmarks[index] def _worker_entry( rank: int, device: str, baseline_code_path: str, extra_args: List[str], benchmarks: List[Tuple[str, Dict[str, Any]]], counter, partial_output_path: str, dump_dir: str, oracle_mode: bool, mg: Optional[str], dump_pred: bool, dump_gt: bool, tqdm_lock, ) -> None: """Entry point of a single-GPU worker process. Must stay at module level to be picklable by `spawn`.""" # Pin this worker to its GPU. This MUST happen before anything imports torch, hence the assertion: # a future module-level `import torch` would otherwise silently put every worker on the same GPU. assert 'torch' not in sys.modules, 'torch was imported before CUDA_VISIBLE_DEVICES could be pinned' os.environ['CUDA_VISIBLE_DEVICES'] = device os.environ.setdefault('OPENCV_IO_ENABLE_OPENEXR', '1') try: from tqdm import tqdm tqdm.set_lock(tqdm_lock) # serialize the progress bars across processes baseline = _load_baseline(baseline_code_path, extra_args) all_metrics = _evaluate_benchmarks( baseline, _claim_benchmarks(benchmarks, counter), metrics_output_path=partial_output_path, dump_dir=dump_dir, oracle_mode=oracle_mode, mg=mg, dump_pred=dump_pred, dump_gt=dump_gt, tqdm_position=rank, tqdm_prefix=f'[gpu {device}] ', ) # Written explicitly instead of relying on the intermediate save, which never fires for a worker # that ends up claiming nothing. Path(partial_output_path).write_text(json.dumps(all_metrics, indent=4)) except KeyboardInterrupt: sys.exit(130) except BaseException: import traceback print(f'\n[rank {rank} | CUDA_VISIBLE_DEVICES={device}] evaluation failed:', file=sys.stderr) traceback.print_exc() sys.stderr.flush() sys.exit(1) def _run_multi_gpu( ngpu: int, baseline_code_path: str, extra_args: Sequence[str], benchmarks: Sequence[Tuple[str, Dict[str, Any]]], output_path: Union[str, Path], dump_dir: Union[str, Path], oracle_mode: bool, mg: Optional[str], dump_pred: bool, dump_gt: bool, ) -> Dict[str, Any]: """Distribute the benchmarks over `ngpu` worker processes and merge their results.""" import multiprocessing as mp devices = _resolve_visible_devices(ngpu) num_workers = min(ngpu, len(benchmarks)) if num_workers < ngpu: click.echo(f'Note: the config has only {len(benchmarks)} benchmark(s); using {num_workers} of the {ngpu} requested GPUs.', err=True) devices = devices[:num_workers] partial_paths = [_partial_output_path(output_path, rank) for rank in range(num_workers)] for path in partial_paths: path.unlink(missing_ok=True) # never merge a leftover file from a previous run # `spawn` (not `fork`, not `torch.multiprocessing.spawn`) gives a fresh interpreter whose first # executed statement is ours, which is what makes the CUDA_VISIBLE_DEVICES pinning above possible. ctx = mp.get_context('spawn') tqdm_lock = ctx.RLock() counter = ctx.Value('i', 0) # index of the next benchmark to be claimed; see `_claim_benchmarks` processes = [] try: for rank in range(num_workers): process = ctx.Process( target=_worker_entry, args=( rank, devices[rank], baseline_code_path, list(extra_args), list(benchmarks), counter, str(partial_paths[rank]), str(dump_dir), oracle_mode, mg, dump_pred, dump_gt, tqdm_lock, ), name=f'eval-rank{rank}-gpu{devices[rank]}', daemon=True, # safe: the workers only ever spawn threads, never processes ) process.start() processes.append(process) for process in processes: process.join() finally: # Leave no orphans behind on Ctrl-C or on a parent-side exception. for process in processes: if process.is_alive(): process.terminate() for process in processes: process.join(timeout=10) if process.is_alive(): process.kill() print('\n' * num_workers, file=sys.stderr, end='') # move the cursor below the pinned progress bars # Collect the per-rank results, treating a crashed worker as a hard failure. all_metrics, failures = {}, [] for rank, (process, partial_path) in enumerate(zip(processes, partial_paths)): if process.exitcode != 0: failures.append(f'rank {rank} (CUDA_VISIBLE_DEVICES={devices[rank]}) exited with code {process.exitcode}') continue if not partial_path.exists(): failures.append(f'rank {rank} (CUDA_VISIBLE_DEVICES={devices[rank]}) produced no result file at {partial_path}') continue all_metrics.update(json.loads(partial_path.read_text())) # Which benchmarks a worker claims is only decided at run time, so completeness is checked globally. # This also catches a worker that died in the middle of a benchmark, leaving a valid but truncated snapshot. if missing := [benchmark_name for benchmark_name, _ in benchmarks if benchmark_name not in all_metrics]: failures.append(f'no result was produced for {missing}') if failures: kept = [path for path in partial_paths if path.exists()] raise click.ClickException( 'the multi-GPU evaluation did not complete:\n' + '\n'.join(f' {failure}' for failure in failures) + f'\n{output_path} was not written.' + ( '\nThe per-rank partial results are kept for inspection:\n' + '\n'.join(f' {path}' for path in kept) if kept else '' ) ) for path in partial_paths: path.unlink(missing_ok=True) return all_metrics @click.command(context_settings={"allow_extra_args": True, "ignore_unknown_options": True}, help='Evaluation script.') @click.option('--baseline', 'baseline_code_path', type=click.Path(), required=True, help='Path to the baseline model python code.') @click.option('--config', 'config_path', type=click.Path(), default='configs/eval/all_benchmarks.json', help='Path to the evaluation configurations. ' 'Defaults to "configs/eval/all_benchmarks.json".') @click.option('--output', '-o', 'output_path', type=click.Path(), required=True, help='Path to the output json file.') @click.option('--ngpu', 'ngpu', type=click.IntRange(min=1), default=1, help='Number of GPUs to use. Whole benchmarks of the config are ' 'distributed over one worker process per GPU, each claiming the next benchmark as it goes. Defaults to 1, i.e. a single in-process run.') @click.option('--oracle', 'oracle_mode', is_flag=True, help='Use oracle mode for evaluation, i.e., use the GT intrinsics input.') @click.option('--mg', 'mg', type=str, default='moge3', help='Comma-separated metric groups to compute.') @click.option('--dump_pred', is_flag=True, help='Dump predition results.') @click.option('--dump_gt', is_flag=True, help='Dump ground truth.') @click.pass_context def main(ctx: click.Context, baseline_code_path: str, config_path: str, ngpu: int, oracle_mode: bool, mg: Optional[str], output_path: Union[str, Path], dump_pred: bool, dump_gt: bool): from moge.utils.tools import key_average # a stdlib-only module: importing it keeps the parent torch-free # Load the evaluation configurations with open(config_path, 'r') as f: config = json.load(f) benchmarks = list(config.items()) if not benchmarks: raise click.UsageError(f'No benchmark is configured in {config_path}.') Path(output_path).parent.mkdir(parents=True, exist_ok=True) dump_dir = Path(str(output_path).replace('.json', '_dump')) if ngpu == 1 or len(benchmarks) == 1: # Single benchmark: the worker count would be capped to 1 anyway, and that worker would get the # very same GPU as an in-process run, so skip the subprocess entirely. baseline = _load_baseline(baseline_code_path, ctx.args) all_metrics = _evaluate_benchmarks( baseline, benchmarks, metrics_output_path=output_path, dump_dir=dump_dir, oracle_mode=oracle_mode, mg=mg, dump_pred=dump_pred, dump_gt=dump_gt, ) else: if any(arg == '--device' or arg.startswith('--device=') for arg in ctx.args): raise click.UsageError( '--device cannot be combined with --ngpu > 1: every worker is pinned to its own GPU via ' 'CUDA_VISIBLE_DEVICES, and an explicit --device would send all of them to the same physical GPU. ' 'Drop --device, or use --ngpu 1.' ) per_benchmark_metrics = _run_multi_gpu( ngpu, baseline_code_path, ctx.args, benchmarks, output_path, dump_dir, oracle_mode, mg, dump_pred, dump_gt, ) # Restore the config order, which depends on which worker happened to claim what. all_metrics = {benchmark_name: per_benchmark_metrics[benchmark_name] for benchmark_name, _ in benchmarks} # Save final results all_metrics['mean'] = key_average(list(all_metrics.values())) Path(output_path).write_text(json.dumps(all_metrics, indent=4)) if __name__ == '__main__': main()