From 47a2753c7e46a7325ce58d2db5374a6131f1c972 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?= Date: Sun, 5 May 2024 09:54:10 -0400 Subject: [PATCH] ENH: Add cerebparc segmentation script fixed DiceCE/Hausdorff weights Add cerebparc segmentation script fixed DiceCE/Hausdorff weights. --- scripts/train_cerebparc_scheduler.py | 636 +++++++++++++++++++++++++++ 1 file changed, 636 insertions(+) create mode 100644 scripts/train_cerebparc_scheduler.py diff --git a/scripts/train_cerebparc_scheduler.py b/scripts/train_cerebparc_scheduler.py new file mode 100644 index 0000000..dcb3856 --- /dev/null +++ b/scripts/train_cerebparc_scheduler.py @@ -0,0 +1,636 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + +import argparse +import logging +import math +import os +import pickle +import shutil +from functools import reduce +from operator import add +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.distributed as dist +from monai.data import ( + CacheDataset, + DataLoader, + decollate_batch, + partition_dataset, +) +from monai.losses import DiceCELoss, HausdorffDTLoss +from monai.metrics import DiceMetric +from monai.utils import set_determinism # first +from monai.utils.enums import LossReduction +from torch.nn.parallel import DistributedDataParallel +from torch.utils.tensorboard import SummaryWriter + +from dmriseg.dataset.utils import ( # get_suit_classnames, + extract_slice, + get_datasets_cerebparc, + get_model, + get_timestamp, + get_transforms, + inference, +) + +# from dmriseg.models.optimization import OptimizationScheduler +from dmriseg.utils import logging_setup +from dmriseg.visualization.plot_utils import ( + boxplot_channel_metric, + get_label_cmap, + plot_loss_and_metric, +) + +set_determinism(seed=0) +torch.backends.cudnn.benchmark = True + +os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:1024" + +logger = logging.getLogger("root") +logger_file_basename = "experiment_logfile.log" + +logging.getLogger("matplotlib.font_manager").disabled = True + + +def _set_up_logger(log_fname): + + logging_setup.set_up(log_fname) + + +def _build_arg_parser(): + + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawTextHelpFormatter + ) + # parser.add_argument( + # "in_train_fa_dirname", + # help="Input training FA image dirname (*.nii.gz)", + # type=Path, + # ) + # parser.add_argument( + # "in_train_md_dirname", + # help="Input training MD image dirname (*.nii.gz)", + # type=Path, + # ) + parser.add_argument( + "in_train_img_dirname", + help="Input training image dirname (*.nii.gz)", + type=Path, + ) + parser.add_argument( + "in_train_labelmap_dirname", + help="Input training data dirname (*.nii.gz)", + type=Path, + ) + parser.add_argument( + "in_valid_img_dirname", + help="Input validation image data dirname (*.nii.gz)", + type=Path, + ) + parser.add_argument( + "in_valid_labelmap_dirname", + help="Input validation labelmap data dirname (*.nii.gz)", + type=Path, + ) + parser.add_argument( + "out_dirname", + help="Output dirname", + type=Path, + ) + return parser + + +def _parse_args(parser): + + args = parser.parse_args() + + return args + + +def main(): + + parser = _build_arg_parser() + args = _parse_args(parser) + + # Set up logger + logger_fname = Path(args.out_dirname).joinpath(logger_file_basename) + _set_up_logger(logger_fname) + + logger.addHandler(logging.StreamHandler()) + + # Parameters + dataset = "cerebellum_cerebparc" + dout = str(args.out_dirname) + meta_pth = os.path.join(dout, "meta_data.pkl") + meta_channel_pkl = os.path.join(dout, "meta_data_channel_metrics.pkl") + model_pth = os.path.join(dout, "model_latest.pth") + meta_loss_weights_pkl = os.path.join(dout, "loss_weights.pkl") + test_2d = False + test_overfit = False + batch_size = 1 + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + # steps = 2*int(1e4) # 2*int(1e6) + use_amp = False # Does not seem stable, best to disable + # validation_steps = 100 # 5000 + spatial_size = (192,) * 3 + patch_size = (192,) * 3 + lr = 1.0e-4 + # Make output folder + if not os.path.isdir(dout): + os.makedirs(dout, exist_ok=True) + elif os.path.isdir(dout) and test_overfit: + shutil.rmtree(dout) + logger.info(dout) + model_pth = None if not os.path.exists(model_pth) else model_pth + log_dir = dout + "/runs" + writer = SummaryWriter(log_dir) + + # DDP + if "LOCAL_RANK" in os.environ: + logger.info("Setting up DDP...", end="") + ddp = True + local_rank = int(os.environ["LOCAL_RANK"]) + # initialize the distributed training process, every GPU runs in a process + dist.init_process_group(backend="nccl", init_method="env://") + device = torch.device(f"cuda:{local_rank}") + num_gpus = dist.get_world_size() + logger.info("done!") + else: + ddp = False + num_gpus = 1 + torch.cuda.set_device(device) + + # Remove background + classnames = ( + "Left_Dentate", + "Right_Dentate", + "left_Interposed", + "Right_Interposed", + "Left_Fastigial", + "Right_Fastigial", + ) # get_suit_classnames()[1:] + + # Get files + datasets = get_datasets_cerebparc( + str(args.in_train_img_dirname), + str(args.in_train_labelmap_dirname), + str(args.in_valid_img_dirname), + str(args.in_valid_labelmap_dirname), + test_overfit=test_overfit, + test_2d=test_2d, + batch_size=batch_size, + ) + train_files = datasets[dataset]["train"] + val_files = datasets[dataset]["val"] + # train_files = train_files[:8] + # val_files = val_files[:8] + num_train = len(train_files) + num_val = len(val_files) + # Set when to run validation + # steps_per_epoch = num_train + validation_interval = 1 # round(validation_steps/steps_per_epoch) + + # Partition per device + if ddp: + train_files = partition_dataset( + data=train_files, + num_partitions=num_gpus, + shuffle=False, + seed=0, + drop_last=False, + even_divisible=True, + )[dist.get_rank()] + val_files = partition_dataset( + data=val_files, + num_partitions=num_gpus, + shuffle=False, + seed=0, + drop_last=False, + even_divisible=True, + )[dist.get_rank()] + + # Get transforms + transforms = get_transforms( + dataset, + target_labels=datasets[dataset]["target_labels"], + device=device, + n_labels=datasets[dataset]["n_labels"], + test_2d=test_2d, + spatial_size=spatial_size, + test_overfit=test_overfit, + patch_size=patch_size, + ) + + # Get data loaders + train_ds = CacheDataset( + data=train_files, + transform=transforms[dataset]["train"], + cache_rate=1.0, + ) + train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True) + + valid_ds = CacheDataset( + data=val_files, + transform=transforms[dataset]["val"], + cache_rate=1.0, + ) + val_loader = DataLoader(valid_ds, batch_size=1, shuffle=False) + + max_epochs = 300 # steps // num_train + + # Get model + model_name = "SegResNet16" + # model_name = "UNet" + out_channels = 7 # 3 # fastigial datasets[dataset]["n_labels"] + 1 + logger.info(f"Number of output channels: {out_channels}") + model = get_model( + model_name, out_channels, device, test_2d=test_2d, model_pth=model_pth + ) + + if ddp: + model = DistributedDataParallel( + model, + device_ids=[device], + find_unused_parameters=False, + # broadcast_buffers=False, + ) + + # Get loss, optimiser and metrics + loss_dice = DiceCELoss( + to_onehot_y=True, + softmax=True, + include_background=False, + smooth_nr=1e-5, + smooth_dr=1e-5, + squared_pred=True, + ) + loss_hausdorff = HausdorffDTLoss( + to_onehot_y=True, + softmax=True, + include_background=False, + alpha=2.0, + sigmoid=False, + other_act=None, + reduction=LossReduction.MEAN, + batch=False, + ) + + optimizer = torch.optim.Adam( + model.parameters(), lr=math.sqrt(batch_size * num_gpus) * lr + ) + # ToDo + # Will probably want reduction="mean_channel" and then do the mean myself + # so that I know which are the best/worst labels at each epoch. Or "none" to + # have a data point per label, per image so that I can plot a boxplot. Have + # two instances if necessary + dice_metric = DiceMetric(include_background=False, reduction="mean") + dice_metric_c = DiceMetric(include_background=False, reduction="none") + dice_metric_batch = DiceMetric( + include_background=False, reduction="mean_batch" + ) + + # Add LR scheduler + # factor = 0.5 + # patience = 10 + # min_lr = 1e-8 + # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + # optimizer, mode="min", factor=factor, patience=patience, min_lr=min_lr + # ) + + # Various + post_pred = transforms["post"] + scaler = torch.cuda.amp.GradScaler() + fig, axs = plt.subplots(3, 3, figsize=(9, 9)) + cmap, norm = get_label_cmap(n_labels=out_channels) + + dim = 0 # sagittal; dim = 1 # coronal; dim = 2 # axial + + # Load/init meta data dict + if os.path.isfile(meta_pth) and model_pth is not None: + with open(meta_pth, "rb") as f: + meta_data = pickle.load(f) + if "epoch" not in meta_data: + meta_data["epoch"] = 0 + if "best_metric" not in meta_data: + meta_data["best_metric"] = -1 + if "loss_values" not in meta_data: + meta_data["loss_values"] = [] + if "metric_values" not in meta_data: + meta_data["metric_values"] = [] + if "best_epoch" not in meta_data: + meta_data["best_epoch"] = 0 + else: + meta_data = dict( + epoch=0, + best_metric=-1, + loss_values=[], + metric_values=[], + best_epoch=0, + ) + + # ToDo + # This can be a little bit too much if we have a lot of classes and a lot + # of validation data samples, as it saves classes x samples epochs. + if os.path.isfile(meta_channel_pkl) and meta_channel_pkl is not None: + with open(meta_channel_pkl, "rb") as f: + meta_data_channel = pickle.load(f) + if "metric_values" not in meta_data_channel: + meta_data_channel["metric_values"] = [] + else: + meta_data_channel = dict( + metric_values=[], + ) + + if ( + os.path.isfile(meta_loss_weights_pkl) + and meta_loss_weights_pkl is not None + ): + with open(meta_loss_weights_pkl, "rb") as f: + meta_data_loss_weights = pickle.load(f) + if "loss_weights" not in meta_data_loss_weights: + meta_data_loss_weights["loss_weights"] = [] + if "to_steal" not in meta_data_loss_weights: + meta_data_loss_weights["to_steal"] = 0 + else: + meta_data_loss_weights = dict( + to_steal=0, + loss_weights=[], + ) + + loss_weights = [0.62, 0.38] # [1, 0.01] + loss_fns = [loss_dice, loss_hausdorff] + to_steal = 0.0015 + # scheduler = OptimizationScheduler(to_steal) + + start_epoch = meta_data["epoch"] + if start_epoch != 0: + # ToDo + # Save the weights in the meta_data dictionary if it is finally found + # out that the Hausdorff loss is helpful and avoid this. + # Compute the expected weights from the start epoch. + # delta_weights = start_epoch * to_steal + loss_weights = [0.62, 0.38] + # loss_weights = [loss_weights[0] - delta_weights, loss_weights[1] + delta_weights] + + # Train/Eval + logger.info("-" * 64) + logger.info("Starting training from epoch {}".format(start_epoch)) + logger.info("-" * 64) + + best_metric = meta_data["best_metric"] + loss_values = meta_data["loss_values"] + metric_values = meta_data["metric_values"] + best_epoch = meta_data["best_epoch"] + + metric_values_channel = meta_data_channel["metric_values"] + meta_loss_weights = meta_data_loss_weights["loss_weights"] + + for epoch in range(start_epoch, max_epochs): + + # Save latest meta data + meta_data["epoch"] = epoch + meta_data["best_metric"] = best_metric + meta_data["loss_values"] = loss_values + meta_data["metric_values"] = metric_values + meta_data["best_epoch"] = best_epoch + with open(meta_pth, "wb") as f: + pickle.dump(meta_data, f) + + meta_data_channel["metric_values"] = metric_values_channel + with open(meta_channel_pkl, "wb") as f: + pickle.dump(meta_data_channel, f) + + meta_data_loss_weights["loss_weights"] = meta_loss_weights + meta_data_loss_weights["to_steal"] = to_steal + with open(meta_loss_weights_pkl, "wb") as f: + pickle.dump(meta_data_loss_weights, f) + + # Train + # ---------------------- + model.train() + epoch_loss = 0 + # For printing a random debug figure + save_fig_ix = np.random.randint(0, len(train_loader)) + + step_epoch = 0 + for ix, batch_data in enumerate(train_loader): + + inputs, labels = ( + batch_data["image"].to(device), + batch_data["label"].to(device), + ) + step_epoch += 1 + + batch_ix = np.random.randint(0, inputs.shape[0]) + + optimizer.zero_grad() + with torch.cuda.amp.autocast(enabled=use_amp): + outputs = model(inputs) + losses = [ + w * loss_fn(outputs, labels) + for loss_fn, w in zip(loss_fns, loss_weights) + ] + loss = reduce(add, losses) + + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + + epoch_loss += loss.item() + + if (epoch % validation_interval == 0) and (save_fig_ix == ix): + with torch.no_grad(): + post_outputs = [ + post_pred(i) for i in decollate_batch(outputs) + ] + axs[0, 0].imshow( + extract_slice(inputs[batch_ix][0], dim=dim).cpu().numpy(), + cmap="gray", + ) + axs[0, 0].set_title("train image") + axs[0, 0].axis("off") + axs[0, 1].imshow( + extract_slice(labels[batch_ix][0], dim=dim).cpu().numpy(), + cmap=cmap, + norm=norm, + interpolation="nearest", + ) + axs[0, 1].set_title("train label") + axs[0, 1].axis("off") + axs[0, 2].imshow( + extract_slice( + post_outputs[batch_ix].argmax(dim=0), dim=dim + ) + .cpu() + .numpy(), + cmap=cmap, + norm=norm, + interpolation="nearest", + ) + axs[0, 2].set_title("train prediction") + axs[0, 2].axis("off") + + epoch_loss /= step_epoch + logger.info( + f"EPOCH={epoch + 1:{' '}{len(str(max_epochs))}}/{max_epochs} |__LOSS (N={num_train})={epoch_loss:.4f} | {get_timestamp()}" + ) + writer.add_scalar("Loss/train", epoch_loss, epoch) + writer.flush() + + loss_values.append(epoch_loss) + + if epoch % validation_interval == 0 or epoch == max_epochs - 1: + # Eval + # ---------------------- + model.eval() + save_fig_ix = np.random.randint(0, len(val_loader)) + + valid_loss = 0 + with torch.no_grad(): + + for ix, batch_data in enumerate(val_loader): + + inputs, labels = ( + batch_data["image"].to(device), + batch_data["label"].to(device), + ) + batch_ix = np.random.randint(0, inputs.shape[0]) + + outputs = inference(inputs, model, use_amp) + losses = [ + w * loss_fn(outputs, labels) + for loss_fn, w in zip(loss_fns, loss_weights) + ] + loss = reduce(add, losses) + valid_loss += loss.item() + + post_outputs = [ + post_pred(i) for i in decollate_batch(outputs) + ] + dice_metric(y_pred=post_outputs, y=labels) + dice_metric_c(y_pred=post_outputs, y=labels) + dice_metric_batch(y_pred=post_outputs, y=labels) + + if save_fig_ix == ix: + axs[1, 0].imshow( + extract_slice(inputs[batch_ix][0], dim=dim) + .cpu() + .numpy(), + cmap="gray", + ) + axs[1, 0].set_title("val image") + axs[1, 0].axis("off") + axs[1, 1].imshow( + extract_slice(labels[batch_ix][0], dim=dim) + .cpu() + .numpy(), + cmap=cmap, + norm=norm, + interpolation="nearest", + ) + axs[1, 1].set_title("val label") + axs[1, 1].axis("off") + axs[1, 2].imshow( + extract_slice( + post_outputs[batch_ix].argmax(dim=0), dim=dim + ) + .cpu() + .numpy(), + cmap=cmap, + norm=norm, + interpolation="nearest", + ) + axs[1, 2].set_title("val prediction") + axs[1, 2].axis("off") + + # aggregate the final mean dice result + metric = dice_metric.aggregate().item() + metric_c = dice_metric_c.aggregate() + metric_batch = dice_metric_batch.aggregate() + # reset the status for next validation round + dice_metric.reset() + dice_metric_c.reset() + dice_metric_batch.reset() + + metric_values.append(metric) + + logger.info( + f"EPOCH={epoch + 1:{' '}{len(str(max_epochs))}}/{max_epochs} |____METRIC (N={num_val})={metric:.4f}" + ) + + writer.add_scalar("Metric/eval", metric, epoch) + writer.flush() + for i in range(0, len(metric_batch), 10): + logger.info( + " " * (13 + len(str(max_epochs))) + + "|______" + + ", ".join( + [ + f"{k + i:3.0f}={v:0.3f}".format(k, v) + for k, v in enumerate(metric_batch[i : i + 10]) + ] + ) + ) + plot_loss_and_metric( + axs, loss_values, metric_values, validation_interval + ) + fig.suptitle( + f"EPOCH={epoch}, LOSS={epoch_loss:.4f}, METRIC={metric:.4f}" + ) + fig.tight_layout() + fig.savefig(os.path.join(dout, "outputs.png")) + + # ToDo + # Save channel metrics over epochs to see how they evolve with + # weights + # ToDo + # Save the best epoch to the meta data as wel + _metric_values_channel = metric_c.cpu().numpy() + metric_values_channel.append(_metric_values_channel) + + title = f"EPOCH={epoch}" + grid = True + _fig = boxplot_channel_metric( + _metric_values_channel, + "Dice", + classnames, + title=title, + grid=grid, + ) + _fig.savefig(os.path.join(dout, "dice_channel_boxplot.png")) + + if metric > best_metric: + best_metric = metric + best_epoch = epoch + torch.save( + model.state_dict(), + os.path.join(dout, "model_best.pth"), + ) + # Save the channel-wise boxplot for the best epoch. + fig.savefig(os.path.join(dout, "best_outputs.png")) + _fig.savefig( + os.path.join(dout, "best_dice_channel_boxplot.png") + ) + + plt.close(fig) + plt.close(_fig) + + # scheduler.step(valid_loss) + # _lr = scheduler._last_lr + # logger.info(f"Learning rate is: {_lr}") + + torch.save(model.state_dict(), os.path.join(dout, "model_latest.pth")) + writer.close() + + meta_loss_weights.append(loss_weights) + # optimizer, loss_fns, loss_weights = scheduler( + # epoch, optimizer, loss_fns, loss_weights + # ) + + +if __name__ == "__main__": + main()