def test_step(self, batch, batch_idx): x, ys = batch for y in ys: ( loss, reconstruction_loss, kl_loss, mu_dist, std_prior, std_posterior, ) = self.punet.elbo(x, y) self.log("test/loss", loss) self.log("test/kl_div", kl_loss) self.log("test/recon_loss", reconstruction_loss) self.log("test/mu_dist", mu_dist) self.log("test/std_post_norm", std_posterior) self.log("test/std_prior_norm", std_prior) if self.hparams.compute_comparison_metrics: for sample_count in [1, 4, 8, 16]: ged, sample_diversity = generalized_energy_distance( self, x, ys, sample_count=sample_count) self.log(f"test/ged/{sample_count}", ged) self.log(f"test/sample_diversity/{sample_count}", sample_diversity) dice = heatmap_dice_loss(self, x, ys, sample_count=sample_count) self.log(f"test/diceloss/{sample_count}", dice)
def validation_step(self, batch, batch_idx): x, ys = batch for y in ys: ( loss, reconstruction_loss, kl_loss, mu_dist, std_prior, std_posterior, ) = self.punet.elbo(x, y) self.log("val/loss", loss) self.log("val/kl_div", kl_loss) self.log("val/recon_loss", reconstruction_loss) self.log("val/mu_dist", mu_dist) self.log("val/std_post_norm", std_posterior) self.log("val/std_prior_norm", std_prior) if self.hparams.compute_comparison_metrics: # calculate aditional metrics every 5 epochs if self.current_epoch % 5 == 0: for sample_count in [1, 4, 8, 16]: ged, sample_diversity = generalized_energy_distance( self, x, ys, sample_count=sample_count) self.log(f"val/ged/{sample_count}", ged) self.log(f"val/sample_diversity/{sample_count}", sample_diversity) dice = heatmap_dice_loss(self, x, ys, sample_count=sample_count) self.log(f"val/diceloss/{sample_count}", dice)
def test_step(self, batch, batch_idx): x, ys = batch y_hats = [model.forward(x) for model in self.models] ensemble_loss = [] # We cycle through constituent models and ground truth annotations, repeating annotations if nessesary for i, (y_hat, y) in enumerate(zip(y_hats, cycle(ys))): loss = self.lossfun(y_hat, y[:, 0]) self.log(f"test/loss_model_{i}", loss) ensemble_loss.append(loss) ensemble_loss = torch.stack(ensemble_loss).mean() self.log("test/loss", ensemble_loss) if self.hparams.compute_comparison_metrics: for sample_count in [1, 4, 8, 16]: ged, sample_diversity = generalized_energy_distance( self, x, ys, sample_count=sample_count ) self.log(f"test/ged/{sample_count}", ged) self.log( f"test/sample_diversity/{sample_count}", sample_diversity) dice = heatmap_dice_loss( self, x, ys, sample_count=sample_count) self.log(f"test/diceloss/{sample_count}", dice)
def test_step(self, batch, batch_idx): x, ys = batch y_hat = self.unet(x) for y in ys: loss = self.lossfun(y_hat, y[:, 0]) self.log("test/loss", loss) if self.hparams.compute_comparison_metrics: for sample_count in [1, 4, 8, 16]: ged, sample_diversity = generalized_energy_distance( self, x, ys, sample_count=sample_count) self.log(f"test/ged/{sample_count}", ged) self.log(f"test/sample_diversity/{sample_count}", sample_diversity) dice = heatmap_dice_loss(self, x, ys, sample_count=sample_count) self.log(f"test/diceloss/{sample_count}", dice)
def validation_step(self, batch, batch_idx): x, ys = batch y_hat = self.unet(x) for y in ys: loss = self.lossfun(y_hat, y[:, 0]) self.log("val/loss", loss) if self.hparams.compute_comparison_metrics: # calculate aditional metrics every 5 epochs if self.current_epoch % 5 == 0: for sample_count in [1, 4, 8, 16]: ged, sample_diversity = generalized_energy_distance( self, x, ys, sample_count=sample_count) self.log(f"val/ged/{sample_count}", ged) self.log(f"val/sample_diversity/{sample_count}", sample_diversity) dice = heatmap_dice_loss(self, x, ys, sample_count=sample_count) self.log(f"val/diceloss/{sample_count}", dice)
def cli_main(): pl.seed_everything(1234) supported_models = util.get_supported_models() # ------------ # args # ------------ parser = ArgumentParser() parser.add_argument( '--model_path', type=str, help=f'Path to the trained model.') parser.add_argument( '--file', type=str, default='./plots/experiment_results.pickl', help=f'File to save the results in.') parser = pl.Trainer.add_argparse_args(parser) args = parser.parse_args() # ------------ # model # ------------ model = util.load_model_from_checkpoint(args.model_path) checkpoint_path = util.get_checkpoint_path(args.model_path) datamodule = util.load_datamodule_for_model(model) dataset = model.hparams.dataset # ------------ # file # ------------ if os.path.isfile(args.file): print('Loading existing result file') with open(args.file, 'rb') as f: test_results = pickle.load(f) else: print('Creating new result file') test_results = {} if dataset not in test_results: test_results[dataset] = {} test_results[dataset][model.model_shortname()] = {} test_results[dataset][model.model_shortname( )]['model_name'] = model.model_name() test_results[dataset][model.model_shortname( )]['model_shortname'] = model.model_shortname() # ------------ # Run model test script # ------------ trainer = pl.Trainer.from_argparse_args(args) test_metrics = trainer.test( model=model, ckpt_path=checkpoint_path, datamodule=datamodule) for k, v in test_metrics[0].items(): test_results[dataset][model.model_shortname()][k] = v # ------------ # Record sample-specific metrics # ------------ with torch.no_grad(): device = 'cuda' if torch.cuda.is_available() else 'cpu' model.to(device) model.eval() metrics = defaultdict(list) pixel_metrics = defaultdict(list) for i, (x, ys) in enumerate(tqdm(datamodule.test_dataloader(), desc='Collecting sample-individual metrics...')): x, ys = util.to_device((x, ys), device) assert ys[0].max() <= 1 y_mean = torch.stack(ys).float().mean(dim=0) # # Image-wise metrics # for sample_count in [1, 4, 8, 16]: ged, sample_diversity = generalized_energy_distance( model, x, ys, sample_count=sample_count) metrics[f"test/ged/{sample_count}"].append(ged) metrics[f"test/sample_diversity/{sample_count}"].append( sample_diversity) dice = heatmap_dice_loss( model, x, ys, sample_count=sample_count) metrics[f"test/diceloss/{sample_count}"].append(dice) sample_count = 16 uncertainty = model.pixel_wise_uncertainty( x, sample_cnt=sample_count) correl = torch.stack([pearsonr(uncertainty, torch.nn.functional.binary_cross_entropy( model.sample_prediction(x).float(), ys[torch.randint(len(ys), ())].float(), reduction='none')) for _ in range(16)]).mean(dim=0) metrics["test/uncertainty_seg_error_correl"].append(correl) # # Pixel-whise metrics # model_uncertainty = uncertainty.cpu().numpy() # model uncertainty values # we sample pixel-wise metrics to preserve memory PIXELS_PER_IMAGE = 1000 B = model_uncertainty.shape[0] indices = np.moveaxis(np.indices(model_uncertainty.shape[1:]), 0, -1) indices = indices.reshape((-1, 3)) indices = indices[np.random.randint(len(indices), size=PIXELS_PER_IMAGE)] # re-add batch dim #indices = np.concatenate([np.stack([np.insert(i, 0, b) for i in indices]) for b in range(B)]) y_hat = model.sample_prediction(x) # model thresholded prediction annotator_sum = torch.stack(ys).sum(dim=0) # sum of annotator votes annotator_cnt = len(ys) model_uncertainty = uncertainty # model uncertainty values y_mean = torch.stack(ys).float().mean(dim=0) annot_uncertainty = util.binary_entropy(y_mean) # annotator uncertainty # record conditional uncetainty only if the is consensus for b in range(B): # iterate over images of the batch # sample pixels by indices idx = torch.tensor([[b]+list(i) for i in indices]) y_hat_subsampled = index_select(y_hat, idx) annotator_sum_subsampled = index_select(annotator_sum, idx) model_uncertainty_subsampled = index_select(model_uncertainty, idx) annot_uncertainty_subsampled = index_select(annot_uncertainty, idx) pixel_metrics["test/tp_uncertainty"].append( model_uncertainty_subsampled[(y_hat_subsampled == 1) & (annotator_sum_subsampled == annotator_cnt)].cpu().tolist()) pixel_metrics["test/fp_uncertainty"].append( model_uncertainty_subsampled[(y_hat_subsampled == 1) & (annotator_sum_subsampled == 0)].cpu().tolist()) pixel_metrics["test/fn_uncertainty"].append( model_uncertainty_subsampled[(y_hat_subsampled == 0) & (annotator_sum_subsampled == annotator_cnt)].cpu().tolist()) pixel_metrics["test/tn_uncertainty"].append( model_uncertainty_subsampled[(y_hat_subsampled == 0) & (annotator_sum_subsampled == 0)].cpu().tolist()) pixel_metrics["test/model_uncertainty"].append(model_uncertainty_subsampled.cpu().tolist()) pixel_metrics["test/annotator_uncertainty"].append(annot_uncertainty_subsampled.cpu().tolist()) pixel_metrics["test/is_prediction_correct"].append((y_hat_subsampled == (annotator_sum_subsampled+2)//4).tolist()) # map metrics into lists of floats test_results[dataset][model.model_shortname()]['per_sample'] = {} for k in metrics: test_results[dataset][model.model_shortname()]['per_sample'][k] = torch.cat( metrics[k]).cpu().numpy() for k in pixel_metrics: test_results[dataset][model.model_shortname()]['per_sample'][k] = pixel_metrics[k] # ------------ # save results # ------------ with open(args.file, 'wb') as f: pickle.dump(test_results, f)