import os import pickle import sys from functools import partial from typing import Callable import numpy as np import pytest import torch from torch.multiprocessing import Pool, set_start_method from pytorch_lightning.metrics import Metric NUM_PROCESSES = 2 NUM_BATCHES = 10 BATCH_SIZE = 16 NUM_CLASSES = 5 EXTRA_DIM = 3 THRESHOLD = 0.5 def setup_ddp(rank, world_size): """ Setup ddp enviroment """ os.environ["MASTER_ADDR"] = 'localhost' os.environ['MASTER_PORT'] = '8088' if torch.distributed.is_available(): torch.distributed.init_process_group("gloo", rank=rank, world_size=world_size) def _class_test( rank: int, worldsize: int, preds: torch.Tensor, target: torch.Tensor, metric_class: Metric, sk_metric: Callable, dist_sync_on_step: bool, metric_args: dict = {}, check_dist_sync_on_step: bool = True, check_batch: bool = True, atol: float = 1e-8, ): """ Utility function doing the actual comparison between lightning class metric and reference metric. Args: rank: rank of current process worldsize: number of processes preds: torch tensor with predictions target: torch tensor with targets metric_class: lightning metric class that should be tested sk_metric: callable function that is used for comparison dist_sync_on_step: bool, if true will synchronize metric state across processes at each ``forward()`` metric_args: dict with additional arguments used for class initialization check_dist_sync_on_step: bool, if true will check if the metric is also correctly calculated per batch per device (and not just at the end) check_batch: bool, if true will check if the metric is also correctly calculated across devices for each batch (and not just at the end) """ # Instanciate lightning metric metric = metric_class(compute_on_step=True, dist_sync_on_step=dist_sync_on_step, **metric_args) # verify metrics work after being loaded from pickled state pickled_metric = pickle.dumps(metric) metric = pickle.loads(pickled_metric) for i in range(rank, NUM_BATCHES, worldsize): batch_result = metric(preds[i], target[i]) if metric.dist_sync_on_step: if rank == 0: ddp_preds = torch.stack([preds[i + r] for r in range(worldsize)]) ddp_target = torch.stack([target[i + r] for r in range(worldsize)]) sk_batch_result = sk_metric(ddp_preds, ddp_target) # assert for dist_sync_on_step if check_dist_sync_on_step: assert np.allclose(batch_result.numpy(), sk_batch_result, atol=atol) else: sk_batch_result = sk_metric(preds[i], target[i]) # assert for batch if check_batch: assert np.allclose(batch_result.numpy(), sk_batch_result, atol=atol) # check on all batches on all ranks result = metric.compute() assert isinstance(result, torch.Tensor) total_preds = torch.stack([preds[i] for i in range(NUM_BATCHES)]) total_target = torch.stack([target[i] for i in range(NUM_BATCHES)]) sk_result = sk_metric(total_preds, total_target) # assert after aggregation assert np.allclose(result.numpy(), sk_result, atol=atol) def _functional_test( preds: torch.Tensor, target: torch.Tensor, metric_functional: Callable, sk_metric: Callable, metric_args: dict = {}, atol: float = 1e-8 ): """ Utility function doing the actual comparison between lightning functional metric and reference metric. Args: preds: torch tensor with predictions target: torch tensor with targets metric_functional: lightning metric functional that should be tested sk_metric: callable function that is used for comparison metric_args: dict with additional arguments used for class initialization """ metric = partial(metric_functional, **metric_args) for i in range(NUM_BATCHES): lightning_result = metric(preds[i], target[i]) sk_result = sk_metric(preds[i], target[i]) # assert its the same assert np.allclose(lightning_result.numpy(), sk_result, atol=atol) class MetricTester: """ Class used for efficiently run alot of parametrized tests in ddp mode. Makes sure that ddp is only setup once and that pool of processes are used for all tests. All tests should subclass from this and implement a new method called `test_metric_name` where the method `self.run_metric_test` is called inside. """ atol = 1e-8 def setup_class(self): """ Setup the metric class. This will spawn the pool of workers that are used for metric testing and setup_ddp """ try: set_start_method('spawn') except RuntimeError: pass self.poolSize = NUM_PROCESSES self.pool = Pool(processes=self.poolSize) self.pool.starmap(setup_ddp, [(rank, self.poolSize) for rank in range(self.poolSize)]) def teardown_class(self): """ Close pool of workers """ self.pool.close() self.pool.join() def run_functional_metric_test( self, preds: torch.Tensor, target: torch.Tensor, metric_functional: Callable, sk_metric: Callable, metric_args: dict = {} ): """ Main method that should be used for testing functions. Call this inside testing method Args: preds: torch tensor with predictions target: torch tensor with targets metric_functional: lightning metric class that should be tested sk_metric: callable function that is used for comparison metric_args: dict with additional arguments used for class initialization """ _functional_test(preds=preds, target=target, metric_functional=metric_functional, sk_metric=sk_metric, metric_args=metric_args, atol=self.atol) def run_class_metric_test( self, ddp: bool, preds: torch.Tensor, target: torch.Tensor, metric_class: Metric, sk_metric: Callable, dist_sync_on_step: bool, metric_args: dict = {}, check_dist_sync_on_step: bool = True, check_batch: bool = True, ): """ Main method that should be used for testing class. Call this inside testing methods. Args: ddp: bool, if running in ddp mode or not preds: torch tensor with predictions target: torch tensor with targets metric_class: lightning metric class that should be tested sk_metric: callable function that is used for comparison dist_sync_on_step: bool, if true will synchronize metric state across processes at each ``forward()`` metric_args: dict with additional arguments used for class initialization check_dist_sync_on_step: bool, if true will check if the metric is also correctly calculated per batch per device (and not just at the end) check_batch: bool, if true will check if the metric is also correctly calculated across devices for each batch (and not just at the end) """ if ddp: if sys.platform == "win32": pytest.skip("DDP not supported on windows") self.pool.starmap( partial( _class_test, preds=preds, target=target, metric_class=metric_class, sk_metric=sk_metric, dist_sync_on_step=dist_sync_on_step, metric_args=metric_args, check_dist_sync_on_step=check_dist_sync_on_step, check_batch=check_batch, atol=self.atol, ), [(rank, self.poolSize) for rank in range(self.poolSize)], ) else: _class_test( 0, 1, preds=preds, target=target, metric_class=metric_class, sk_metric=sk_metric, dist_sync_on_step=dist_sync_on_step, metric_args=metric_args, check_dist_sync_on_step=check_dist_sync_on_step, check_batch=check_batch, atol=self.atol, )