def main(): parser = argparse.ArgumentParser() parser.add_argument('-modelType', default=4, type=int, help='Refer train_utils.py ') parser.add_argument('-numSpkrs', default=7323, type=int, help='Number of output labels for model') parser.add_argument('modelDirectory', help='Directory containing the model checkpoints') parser.add_argument( 'featDir', help='Directory containing features ready for extraction') parser.add_argument('embeddingDir', help='Output directory') args = parser.parse_args() modelFile = max(glob.glob(args.modelDirectory + '/*'), key=os.path.getctime) # Load model definition if args.modelType == 3: net = simpleTDNN(args.numSpkrs, p_dropout=0) else: net = xvecTDNN(args.numSpkrs, p_dropout=0) checkpoint = torch.load(modelFile, map_location=torch.device('cuda')) new_state_dict = OrderedDict() for k, v in checkpoint['model_state_dict'].items(): if k.startswith('module.'): new_state_dict[k[7:]] = v # ugly fix to remove 'module' from key else: new_state_dict[k] = v # load trained weights net.load_state_dict(new_state_dict) net = net.cuda() net.eval() # Parallel Processing try: nSplits = int( sorted(glob.glob(args.featDir + '/split*'), key=getSplitNum)[-1].split('/')[-1].lstrip('split')) except: print('Cannot find %s/splitN directory' % args.featDir) sys.exit(1) if not os.path.isdir(args.embeddingDir): os.makedirs(args.embeddingDir) nProcs = nSplits L = [('%s/split%d/%d/feats.scp' % (args.featDir, nSplits, i), '%s/xvector.%d.ark' % (args.embeddingDir, i), '%s/xvector.%d.scp' % (args.embeddingDir, i), net, 'fc1') for i in range(1, nSplits + 1)] pool2 = Pool(processes=nProcs) result = pool2.starmap(par_core_extractXvectors, L) pool2.terminate() os.system('cat %s/xvector.*.scp > %s/xvector.scp' % (args.embeddingDir, args.embeddingDir))
def meta_ars(env_name, policy, meta_epochs, meta_seed, n_seeds=4, n_top_seeds=1, n_workers=4, mean_lookback=10, ars_epochs=10, env_config=None, step_size=.02, n_delta=32, n_top=16, exp_noise=0.03): n_children = n_seeds // n_top_seeds np.random.seed(meta_seed) W = torch.nn.utils.parameters_to_vector(policy.parameters()) W = torch.zeros_like(W) torch.nn.utils.vector_to_parameters(W, policy.parameters()) pool = Pool(processes=n_seeds) ars_partial = partial(ars, env_name, ars_epochs, env_config, step_size, n_delta, n_top, exp_noise, n_workers) #root = Node(meta_seed) reward_log = [] top_policies = [] for _ in range(n_top_seeds): top_policies.append(copy.deepcopy(policy)) for epoch in range(meta_epochs): pols_and_seeds = [] for pol in top_policies: for _ in range(n_children): pols_and_seeds.append( (pol, int(np.random.randint(0, 2**32 - 1, 1)))) results = pool.starmap(ars_partial, pols_and_seeds) p_list = [] r_list = [] for result in results: policy, rews = result p_list.append(policy) r = torch.stack(rews[-mean_lookback:]) r_list.append(r.mean()) top_idx = sorted(range(len(r_list)), key=lambda k: r_list[k], reverse=True)[:n_top_seeds] for i in top_idx: top_policies.append(p_list[i]) reward_log.append(max(r_list)) return top_policies, reward_log
class Predictor: def __init__(self, cfg_path: str, num_workers: int = 4) -> None: cfg = get_cfg() cfg.merge_from_file(model_zoo.get_config_file(cfg_path)) # NOTE: you may customize cfg settings # cfg.MODEL.DEVICE="cuda" # use gpu by default cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well # you can also give a path to you checkpoint cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(cfg_path) self.cfg = cfg.clone() self.model = build_model(cfg) self.model.eval() checkpointer = DetectionCheckpointer(self.model) checkpointer.load(cfg.MODEL.WEIGHTS) self.aug = T.ResizeShortestEdge( [cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST) self.pool = Pool(num_workers) @staticmethod def url_to_img(aug: object, device: str, url: str) -> Dict: img_response = requests.get(url) img = np.array(Image.open(BytesIO(img_response.content))) height, width = img.shape[:2] img = aug.get_transform(img).apply_image(img) img = torch.as_tensor(img.astype("float32").transpose(2, 0, 1)) if device == "cuda": img = img.pin_memory() img = img.cuda(non_blocking=True) return {"image": img, "height": height, "width": width} def load_inputs(self, urls: List[str]) -> List[Dict]: # return [self.url_to_img(url) for url in urls] n = len(urls) return list( self.pool.starmap( self.url_to_img, zip([self.aug] * n, [self.cfg.MODEL.DEVICE] * n, urls))) def __call__(self, items: Dict) -> List: urls = [item["url"] for item in items] inputs = self.load_inputs(urls) with torch.no_grad(): predictions = self.model(inputs) return predictions
def create_episodes( self, n_episodes: int, n_processes: int, mcts_samples: int, mcts_temp: float, mcts_cpuct: int, mcts_observation_weight: float, model: Model, ) -> List[Tuple[List[ObservationType], List[np.ndarray], int, Summary]]: pool = Pool(n_processes) res = pool.starmap( self._generator.perform_episode, [[mcts_samples, mcts_temp, mcts_cpuct, mcts_observation_weight, model]] * n_episodes, ) pool.close() pool.terminate() pool.join() return res
def distribute_action_among_workers( pool: mp.Pool, num_workers: int, config: ConfigSchema, action: Action, model: MultiRelationEmbedder, epoch_idx: int, lhs: EntityList, rhs: EntityList, rel: LongTensorType, edge_perm: LongTensorType, optimizers: Optional[List[Optimizer]] = None ) -> Union[TrainStats, EvalStats]: all_stats = pool.starmap(perform_action_one_thread, [ (Rank(i), config, action, model, epoch_idx, lhs, rhs, rel, edge_perm[s], optimizers) for i, s in enumerate(split_almost_equally(edge_perm.size(0), num_parts=num_workers)) ]) if action is action.TRAIN: return TrainStats.sum(all_stats).average() elif action is action.EVAL: return EvalStats.sum(all_stats).average() else: raise NotImplementedError("Unknown action: %s" % action)
def start_multiprocessing(self, embeddings, naming_list, naming_dict, dataset): n_threads = 4 chunked = chunkIt(naming_list, n_threads) # multiprocess gridsearch and have a seperate thread for the progress bar. pool1 = Pool(processes=n_threads) m = Manager() q = m.Queue() p = Process(target=progressBar, args=(len(naming_list), q,)) p.start() results = pool1.starmap(self.determine_triplets, zip(n_threads * [q], chunked, n_threads * [naming_list], n_threads * [embeddings], n_threads * [naming_dict], n_threads * [dataset])) final_results = [] for r in results: final_results += r p.join() pool1.close() return final_results
def main(): parser = argparse.ArgumentParser() parser.add_argument('-modelType', default='xvecTDNN', help='Refer train_utils.py ') parser.add_argument('-numSpkrs', default=7323, type=int, help='Number of output labels for model') parser.add_argument('-layerName', default='fc1', help="DNN layer for embeddings") parser.add_argument('modelDirectory', help='Directory containing the model checkpoints') parser.add_argument( 'featDir', help='Directory containing features ready for extraction') parser.add_argument('embeddingDir', help='Output directory') args = parser.parse_args() try: modelFile = max(glob.glob(args.modelDirectory + '/*'), key=os.path.getctime) except ValueError: print("[ERROR] No trained model has been found in {}.".format( args.modelDirectory)) sys.exit(1) # Load model definition net = eval('{}({}, p_dropout=0)'.format(args.modelType, args.numSpkrs)) checkpoint = torch.load(modelFile, map_location=torch.device('cuda')) new_state_dict = OrderedDict() if 'relation' in args.modelType: checkpoint_dict = checkpoint['encoder_state_dict'] else: checkpoint_dict = checkpoint['model_state_dict'] for k, v in checkpoint_dict.items(): if k.startswith('module.'): new_state_dict[k[7:]] = v # ugly fix to remove 'module' from key else: new_state_dict[k] = v # load trained weights net.load_state_dict(new_state_dict) net = net.cuda() net.eval() # Parallel Processing try: nSplits = int( sorted(glob.glob(args.featDir + '/split*'), key=getSplitNum)[-1].split('/')[-1].lstrip('split')) except ValueError: print('[ERROR] Cannot find %s/splitN directory' % args.featDir) sys.exit(1) if not os.path.isdir(args.embeddingDir): os.makedirs(args.embeddingDir) print('Extracting xvectors by distributing jobs to pool workers... ') nProcs = nSplits L = [('%s/split%d/%d/feats.scp' % (args.featDir, nSplits, i), '%s/xvector.%d.ark' % (args.embeddingDir, i), '%s/xvector.%d.scp' % (args.embeddingDir, i), net, args.layerName) for i in range(1, nSplits + 1)] pool2 = Pool(processes=nProcs) result = pool2.starmap(par_core_extractXvectors, L) pool2.terminate() print('Multithread job has been finished.') print('Writing xvectors to {}'.format(args.embeddingDir)) os.system('cat %s/xvector.*.scp > %s/xvector.scp' % (args.embeddingDir, args.embeddingDir))
def eval_(self, X, *args): """ Evaluate a number of DARTS architecture in parallel. X should be a list of Genotypes defined by DARTS API. """ from math import ceil n_parallel = min(len(X), self.n_gpu) res = [] diag_stats = [] if n_parallel == 0: raise ValueError("No GPUs available!") elif n_parallel == 1: for i, genotype in enumerate(X): t = DARTSTrainer(self.data_path, self.save_path, genotype, self.dataset, cutout=self.cutout, auxiliary_tower=self.auxiliary, epochs=self.epochs, eval_policy=self.query_policy) print('Start training: ', i + 1, "/ ", len(X)) try: t.train() # bottleneck result = t.retrieve() res.append(1. - result[0] / 100.) # Turn into error diag_stats.append(result[1]) except Exception as e: logging.error( "An error occured in the current architecture. Assigning nan to the arch. The error is:" ) logging.error(e) res.append(np.nan) diag_stats.append(None) else: gpu_ids = range(n_parallel) num_reps = ceil(len(X) / float(n_parallel)) for i in range(num_reps): x = X[i * n_parallel:min((i + 1) * n_parallel, len( X))] # select the number of parallel archs to evaluate selected_gpus = gpu_ids[:len(x)] other_arg = [ self.data_path, self.save_path, self.dataset, self.cutout, self.epochs, self.query_policy ] args = list(map( list, zip( x, selected_gpus, ), )) args = [a + other_arg for a in args] pool = Pool(processes=len(x)) current_res = pool.starmap(parallel_eval, args) pool.close() pool.join() res.extend([i for i in current_res if i >= 0 ]) # Filter out the negative results due to errors res = np.array(res).flatten() if self.log_scale: res = np.log(res) if self.negative: res = -res return res, diag_stats
# print(probabilities) break return (features, payoffs, probabilities) game = GameState(4) features = [] payoffs = [] probabilities = [] policy = None step = 128 milestone = step mcts_pool = Pool(32) for x in range(100): results: List[Tuple[Any, Any, Any]] = mcts_pool.starmap(one_mcts, [(game, policy)] * 64) # f, pa, pr = one_mcts(game, policy) for f, pa, pr in results: features.extend(f) payoffs.extend(pa) probabilities.extend(pr) while len(features) >= milestone: milestone += step if policy is None: policy = Policy(game.feature_dim(), game.action_dim(), game.num_players) policy.eval() policy.train() policy.fit(torch.stack(features), torch.stack(payoffs), torch.stack(probabilities))
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 """ 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: Tensor, target: Tensor, metric_functional: Callable, sk_metric: Callable, metric_args: dict = None, **kwargs_update, ): """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 kwargs_update: Additional keyword arguments that will be passed with preds and target when running update on the metric. """ _functional_test( preds=preds, target=target, metric_functional=metric_functional, sk_metric=sk_metric, metric_args=metric_args, atol=self.atol, **kwargs_update, ) def run_class_metric_test( self, ddp: bool, preds: Tensor, target: Tensor, metric_class: Metric, sk_metric: Callable, dist_sync_on_step: bool, metric_args: dict = None, check_dist_sync_on_step: bool = True, check_batch: bool = True, **kwargs_update, ): """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) kwargs_update: Additional keyword arguments that will be passed with preds and target when running update on the metric. """ if not metric_args: metric_args = {} 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, **kwargs_update, ), [(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, **kwargs_update, ) def run_precision_test_cpu( self, preds: torch.Tensor, target: torch.Tensor, metric_module: Metric, metric_functional: Callable, metric_args: dict = {}, ): """Test if an metric can be used with half precision tensors on cpu Args: preds: torch tensor with predictions target: torch tensor with targets metric_module: the metric module to test metric_functional: the metric functional to test metric_args: dict with additional arguments used for class initialization """ _assert_half_support(metric_module(**metric_args), partial(metric_functional, **metric_args), preds, target, device="cpu") def run_precision_test_gpu( self, preds: torch.Tensor, target: torch.Tensor, metric_module: Metric, metric_functional: Callable, metric_args: dict = {}, ): """Test if an metric can be used with half precision tensors on gpu Args: preds: torch tensor with predictions target: torch tensor with targets metric_module: the metric module to test metric_functional: the metric functional to test metric_args: dict with additional arguments used for class initialization """ _assert_half_support(metric_module(**metric_args), partial(metric_functional, **metric_args), preds, target, device="cuda")
class ESOptimizer: """ An optimizer class that implements Evolution Strategies (ES) """ def __init__(self, model: Model, sgd_optimizer: Optimizer, objective_fn: Objective, obj_weights: List[float], sigma: float, n_samples: int, devices: List, n_workers=4): self.model = model self._optimizer = sgd_optimizer self.sigma = sigma self.n_samples = n_samples self.objective_fn = objective_fn self.obj_weights = torch.Tensor(obj_weights) self.n_objectives = len(obj_weights) self.devices = devices self.pool = Pool(processes=n_workers) # evaluator self.evaluator = ModelEvaluator(self.model, self.objective_fn, None, None) def _compute_gradient(self, obj_value: List[float], delta: float): """ Computes the gradient for one sample """ obj_value = torch.Tensor(obj_value) weighted_sum = torch.dot(obj_value, self.obj_weights) grad = delta * weighted_sum return grad def _fit_scalers_for_objectives( self, objectives: List[List[float]]) -> List[RankScaler]: """ Fits rank scalers for each objectives """ rank_scalers = [] n_values = len(objectives) for obj_ix in range(self.n_objectives): values = [objectives[i][obj_ix] for i in range(n_values)] scaler = RankScaler(values) rank_scalers.append(scaler) return rank_scalers def _gradients_from_objectives( self, current_value: torch.Tensor, obj_values: List[List[float]], perturbations: List[torch.Tensor]) -> torch.Tensor: """ Computes average gradients using multi-objective values """ total_gradients = torch.zeros(current_value.shape) rank_scalers = self._fit_scalers_for_objectives(obj_values) for obj, delta in zip(obj_values, perturbations): # rank scale them obj = [ rank_scalers[ix].transform(value) for ix, value in enumerate(obj) ] # compute gradient gradient = self._compute_gradient(obj, delta) total_gradients += gradient # average the gradients grad = total_gradients / (self.n_samples * self.sigma) return grad def _generate_perturbations(self, current_value): # create mirrored sampled perturbations perturbs = [ torch.randn_like(current_value) for i in range(int(self.n_samples / 2)) ] mirrored = [-i for i in perturbs] perturbs += mirrored return perturbs def gradient_step(self, samples: Samples): """ Performs a gradient ascent step Args: samples (Samples): samples """ # sample some parameters here parameter_name, parameter_value = self.model.sample() # generate unit gaussian perturbations unit_perturbations = self._generate_perturbations(parameter_value) # apply user selected deviation perturbations = [ parameter_value + perturb * self.sigma for perturb in unit_perturbations ] # sample devices devices = random.choices(self.devices, k=len(unit_perturbations)) # get the objective values self.evaluator.current_parameter_name = parameter_name self.evaluator.samples = samples obj_values = self.pool.starmap(self.evaluator, zip(perturbations, devices)) # compute gradients gradients = self._gradients_from_objectives(parameter_value, obj_values, unit_perturbations) # update the model paramters and take a gradient step self.model.set_gradients(parameter_name, -gradients) self._optimizer.step() return gradients
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 """ 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: Tensor, target: Tensor, metric_functional: Callable, sk_metric: Callable, metric_args: dict = None, fragment_kwargs: bool = False, **kwargs_update, ): """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 fragment_kwargs: whether tensors in kwargs should be divided as `preds` and `target` among processes kwargs_update: Additional keyword arguments that will be passed with preds and target when running update on the metric. """ device = 'cuda' if (torch.cuda.is_available() and torch.cuda.device_count() > 0) else 'cpu' _functional_test( preds=preds, target=target, metric_functional=metric_functional, sk_metric=sk_metric, metric_args=metric_args, atol=self.atol, device=device, fragment_kwargs=fragment_kwargs, **kwargs_update, ) def run_class_metric_test( self, ddp: bool, preds: Tensor, target: Tensor, metric_class: Metric, sk_metric: Callable, dist_sync_on_step: bool, metric_args: dict = None, check_dist_sync_on_step: bool = True, check_batch: bool = True, fragment_kwargs: bool = False, check_scriptable: bool = True, **kwargs_update, ): """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) fragment_kwargs: whether tensors in kwargs should be divided as `preds` and `target` among processes kwargs_update: Additional keyword arguments that will be passed with preds and target when running update on the metric. """ if not metric_args: metric_args = {} 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, fragment_kwargs=fragment_kwargs, check_scriptable=check_scriptable, **kwargs_update, ), [(rank, self.poolSize) for rank in range(self.poolSize)], ) else: device = 'cuda' if (torch.cuda.is_available() and torch.cuda.device_count() > 0) else 'cpu' _class_test( rank=0, worldsize=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, device=device, fragment_kwargs=fragment_kwargs, check_scriptable=check_scriptable, **kwargs_update, ) def run_precision_test_cpu( self, preds: Tensor, target: Tensor, metric_module: Metric, metric_functional: Callable, metric_args: dict = None, **kwargs_update, ): """Test if a metric can be used with half precision tensors on cpu Args: preds: torch tensor with predictions target: torch tensor with targets metric_module: the metric module to test metric_functional: the metric functional to test metric_args: dict with additional arguments used for class initialization kwargs_update: Additional keyword arguments that will be passed with preds and target when running update on the metric. """ metric_args = metric_args or {} _assert_half_support(metric_module(**metric_args), metric_functional, preds, target, device="cpu", **kwargs_update) def run_precision_test_gpu( self, preds: Tensor, target: Tensor, metric_module: Metric, metric_functional: Callable, metric_args: dict = None, **kwargs_update, ): """Test if a metric can be used with half precision tensors on gpu Args: preds: torch tensor with predictions target: torch tensor with targets metric_module: the metric module to test metric_functional: the metric functional to test metric_args: dict with additional arguments used for class initialization kwargs_update: Additional keyword arguments that will be passed with preds and target when running update on the metric. """ metric_args = metric_args or {} _assert_half_support(metric_module(**metric_args), metric_functional, preds, target, device="cuda", **kwargs_update) def run_differentiability_test( self, preds: Tensor, target: Tensor, metric_module: Metric, metric_functional: Callable, metric_args: dict = None, ): """Test if a metric is differentiable or not Args: preds: torch tensor with predictions target: torch tensor with targets metric_module: the metric module to test metric_args: dict with additional arguments used for class initialization """ metric_args = metric_args or {} # only floating point tensors can require grad metric = metric_module(**metric_args) if preds.is_floating_point(): preds.requires_grad = True out = metric(preds[0], target[0]) # metrics can return list of values if isinstance(out, list): assert all(metric.is_differentiable == o.requires_grad for o in out) else: assert metric.is_differentiable == out.requires_grad if metric.is_differentiable: # check for numerical correctness assert torch.autograd.gradcheck( partial(metric_functional, **metric_args), (preds[0].double(), target[0])) # reset as else it will carry over to other tests preds.requires_grad = False
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. """ 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_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. 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( _compute_batch, 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, ), [(rank, self.poolSize) for rank in range(self.poolSize)], ) else: _compute_batch( 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, )
class Boundaryloss(object): def __init__(self, dense_crf: DenseCRF, eps: float = 1e-5, crf_num_workers: int = 4): """Compute boundary loss Args: dense_crf: DenseCRF functor eps: Min prob allowed when clamp probs crf_num_workers: num of workers when crf parallel """ self.dense_crf = dense_crf self.eps = eps self.crf_num_workers = crf_num_workers self.crf_pool = Pool(self.crf_num_workers) def crf(self, imgs, probs): np_imgs = BGR2RGB(imgs.cpu().numpy().astype(np.uint8).transpose( 0, 2, 3, 1)) # (N, H, W, C) np_probs = probs.detach().cpu().numpy() # (N, C, H, W) # Scaled imgs to probs shape scaled_imgs = nd.zoom(np_imgs, (1.0, np_probs.shape[2] / np_imgs.shape[1], np_probs.shape[3] / np_imgs.shape[2], 1.0), order=1) # CRF crf_probs = self.crf_pool.starmap(self.dense_crf, zip(scaled_imgs, np_probs)) crf_prob = np.stack(crf_probs, axis=0) # Clamp smoothed probs # TODO: Can be removed? crf_prob[crf_prob < self.eps] = self.eps crf_prob = crf_prob / np.sum(crf_prob, axis=1, keepdims=True) # to Tensor return torch.from_numpy(crf_prob).float().cuda(probs.get_device()) def clamp_softmax(self, score, dim=1): probs = torch.clamp(F.softmax(score, dim), self.eps, 1) probs = probs / torch.sum(probs, dim=dim, keepdim=True) return probs def __call__( self, images, score_map, out_prob=False ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: """Compute the constrain-to-boundary loss Args: images: (N, 3, H, W) RGB img score_map: (N, C, H, W) score map out_prob: If true, return smoothed predict_probs. (Default: False) Returns: constrain-to-boundary loss """ probs = self.clamp_softmax(score_map) smooth_probs = self.crf(images, probs) # Compute KL-Div # TODO: clamp is not needed? loss = torch.mean( torch.sum(smooth_probs * torch.log(torch.clamp(smooth_probs / probs, 0.05, 20)), dim=1)) if out_prob: return loss, smooth_probs return loss def __del__(self): self.crf_pool.close() self.crf_pool.join()