hausdorff_loss_timings.ipynb

{"cells":[{"cell_type":"markdown","metadata":{"id":"V5dCXTLaXB3n"},"source":["# Overview\n","\n","This notebook compares different loss functions for training a 3D segmentation model using MONAI. We focus on three loss function combinations:\n","\n","1. Pure Dice Loss (Gold Standard)\n","2. PyDTs HD Loss + Dice Loss (Proposed method)\n","3. Scipy HD Loss + Dice Loss\n","\n","The notebook has three main sections:\n","\n","**Loss Function Timings**\n","\n","We benchmark the computation timings for each loss function:\n","1. Pure Dice Loss\n","2. PyDTs HD Loss\n","3. Scipy HD Loss\n","\n","outside of the training loop and in isolation. This helps us understand the computational efficiency of each loss function independently. Results are saved to CSV files.\n","\n","**Training Loop Timings**\n","\n","We benchmark the training loop timings for each loss function combination:\n","\n","1. Pure Dice Loss\n","2. PyDTs HD Loss + Dice Loss (proposed method)\n","3. Scipy HD Loss + Dice Loss\n","\n","We run the training loop for 10 epochs and measure the total training time, average epoch time, and standard deviation of epoch time. This helps us assess the computational efficiency of each loss function combination within the context of a training loop. Results are saved to CSV files."]},{"cell_type":"markdown","source":["## Setup"],"metadata":{"id":"dVjrQK_e3bdX"}},{"cell_type":"code","source":["from google.colab import drive\n","drive.mount('/content/drive')"],"metadata":{"colab":{"base_uri":"https://localhost:8080/","height":311},"id":"2-C7anO4qZs3","executionInfo":{"status":"error","timestamp":1720030445402,"user_tz":420,"elapsed":4746,"user":{"displayName":"Dale Black","userId":"06644162883191190715"}},"outputId":"801a1431-af4a-4122-c618-5b528fdb317c"},"execution_count":null,"outputs":[{"output_type":"error","ename":"MessageError","evalue":"Error: credential propagation was unsuccessful","traceback":["\u001b[0;31m---------------------------------------------------------------------------\u001b[0m","\u001b[0;31mMessageError\u001b[0m                              Traceback (most recent call last)","\u001b[0;32m<ipython-input-1-d5df0069828e>\u001b[0m in \u001b[0;36m<cell line: 2>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mgoogle\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcolab\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mdrive\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mdrive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmount\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'/content/drive'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m","\u001b[0;32m/usr/local/lib/python3.10/dist-packages/google/colab/drive.py\u001b[0m in \u001b[0;36mmount\u001b[0;34m(mountpoint, force_remount, timeout_ms, readonly)\u001b[0m\n\u001b[1;32m     98\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mmount\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmountpoint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mforce_remount\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtimeout_ms\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m120000\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mreadonly\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     99\u001b[0m   \u001b[0;34m\"\"\"Mount your Google Drive at the specified mountpoint path.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 100\u001b[0;31m   return _mount(\n\u001b[0m\u001b[1;32m    101\u001b[0m       \u001b[0mmountpoint\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    102\u001b[0m       \u001b[0mforce_remount\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mforce_remount\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.10/dist-packages/google/colab/drive.py\u001b[0m in \u001b[0;36m_mount\u001b[0;34m(mountpoint, force_remount, timeout_ms, ephemeral, readonly)\u001b[0m\n\u001b[1;32m    131\u001b[0m   )\n\u001b[1;32m    132\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mephemeral\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 133\u001b[0;31m     _message.blocking_request(\n\u001b[0m\u001b[1;32m    134\u001b[0m         \u001b[0;34m'request_auth'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    135\u001b[0m         \u001b[0mrequest\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m{\u001b[0m\u001b[0;34m'authType'\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'dfs_ephemeral'\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;32m/usr/local/lib/python3.10/dist-packages/google/colab/_message.py\u001b[0m in \u001b[0;36mblocking_request\u001b[0;34m(request_type, request, timeout_sec, parent)\u001b[0m\n\u001b[1;32m    174\u001b[0m       \u001b[0mrequest_type\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrequest\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparent\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mparent\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mexpect_reply\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    175\u001b[0m   )\n\u001b[0;32m--> 176\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mread_reply_from_input\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrequest_id\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtimeout_sec\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m","\u001b[0;32m/usr/local/lib/python3.10/dist-packages/google/colab/_message.py\u001b[0m in \u001b[0;36mread_reply_from_input\u001b[0;34m(message_id, timeout_sec)\u001b[0m\n\u001b[1;32m    101\u001b[0m     ):\n\u001b[1;32m    102\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0;34m'error'\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mreply\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 103\u001b[0;31m         \u001b[0;32mraise\u001b[0m \u001b[0mMessageError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mreply\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'error'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    104\u001b[0m       \u001b[0;32mreturn\u001b[0m \u001b[0mreply\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'data'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    105\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n","\u001b[0;31mMessageError\u001b[0m: Error: credential propagation was unsuccessful"]}]},{"cell_type":"code","source":["!pip install py_distance_transforms\n","!python -c \"import monai\" || pip install -q \"monai-weekly[gdown, nibabel, tqdm, ignite]\"\n","!python -c \"import matplotlib\" || pip install -q matplotlib\n","%matplotlib inline"],"metadata":{"id":"hDK-XZUlXzlW","collapsed":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["**NOTE**: *First time importing `py_distance_transforms` might take a while (~up to 8 mins)*"],"metadata":{"id":"F7MqAbZu7kZp"}},{"cell_type":"code","execution_count":null,"metadata":{"tags":[],"id":"jxw9rW05XB3p"},"outputs":[],"source":["from py_distance_transforms import transform_cuda\n","from monai.utils import first, set_determinism\n","from monai.transforms import (\n","    AsDiscrete,\n","    AsDiscreted,\n","    EnsureChannelFirstd,\n","    Compose,\n","    CropForegroundd,\n","    LoadImaged,\n","    Orientationd,\n","    RandCropByPosNegLabeld,\n","    SaveImaged,\n","    ScaleIntensityRanged,\n","    Spacingd,\n","    Invertd,\n",")\n","from monai.handlers.utils import from_engine\n","from monai.networks.nets import UNet\n","from monai.networks.layers import Norm\n","from monai.metrics import DiceMetric, HausdorffDistanceMetric, compute_percent_hausdorff_distance, compute_iou, MeanIoU\n","from monai.losses import DiceLoss\n","from monai.inferers import sliding_window_inference\n","from monai.data import CacheDataset, DataLoader, Dataset, decollate_batch\n","from monai.config import print_config\n","from monai.apps import download_and_extract\n","import torch\n","import matplotlib.pyplot as plt\n","import tempfile\n","import shutil\n","import os\n","import glob\n","\n","from scipy.ndimage import distance_transform_edt\n","import torch.nn.functional as F\n","import numpy as np\n","import time\n","import timeit\n","import pandas as pd\n","\n","from juliacall import Main as jl\n","jl.seval(\"import CUDA\")\n","\n","# print_config()"]},{"cell_type":"markdown","metadata":{"id":"ibARsdT3XB3q"},"source":["**Setup data directory**\n","\n","You can specify a directory with the `MONAI_DATA_DIRECTORY` environment variable.  \n","This allows you to save results and reuse downloads.  \n","If not specified a temporary directory will be used."]},{"cell_type":"code","execution_count":null,"metadata":{"tags":[],"id":"XBWMZ3mkXB3q"},"outputs":[],"source":["directory = os.environ.get(\"MONAI_DATA_DIRECTORY\")\n","root_dir = tempfile.mkdtemp() if directory is None else directory\n","print(root_dir)"]},{"cell_type":"markdown","metadata":{"id":"CMc-JgqJXB3q"},"source":["**Download dataset**\n","\n","Downloads and extracts the dataset.  \n","The dataset comes from http://medicaldecathlon.com/."]},{"cell_type":"code","execution_count":null,"metadata":{"tags":[],"id":"FhyT06dDXB3q"},"outputs":[],"source":["resource = \"https://msd-for-monai.s3-us-west-2.amazonaws.com/Task09_Spleen.tar\"\n","md5 = \"410d4a301da4e5b2f6f86ec3ddba524e\"\n","\n","compressed_file = os.path.join(root_dir, \"Task09_Spleen.tar\")\n","data_dir = os.path.join(root_dir, \"Task09_Spleen\")\n","if not os.path.exists(data_dir):\n","    download_and_extract(resource, compressed_file, root_dir, md5)"]},{"cell_type":"markdown","metadata":{"id":"SHBmxnUeXB3q"},"source":["**Set MSD Spleen dataset path**"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"8XXNvy0BXB3r"},"outputs":[],"source":["train_images = sorted(glob.glob(os.path.join(data_dir, \"imagesTr\", \"*.nii.gz\")))\n","train_labels = sorted(glob.glob(os.path.join(data_dir, \"labelsTr\", \"*.nii.gz\")))\n","data_dicts = [{\"image\": image_name, \"label\": label_name} for image_name, label_name in zip(train_images, train_labels)]\n","train_files, val_files = data_dicts[:-9], data_dicts[-9:]"]},{"cell_type":"markdown","metadata":{"id":"0NkGlkfTXB3r"},"source":["**Set deterministic training for reproducibility**"]},{"cell_type":"code","execution_count":null,"metadata":{"id":"2ns_Wp7sXB3r"},"outputs":[],"source":["set_determinism(seed=0)"]},{"cell_type":"markdown","metadata":{"id":"j9_lquWaXB3r"},"source":["**Setup transforms for training and validation**\n","\n","Here we use several transforms to augment the dataset:\n","1. `LoadImaged` loads the spleen CT images and labels from NIfTI format files.\n","1. `EnsureChannelFirstd` ensures the original data to construct \"channel first\" shape.\n","1. `Orientationd` unifies the data orientation based on the affine matrix.\n","1. `Spacingd` adjusts the spacing by `pixdim=(1.5, 1.5, 2.)` based on the affine matrix.\n","1. `ScaleIntensityRanged` extracts intensity range [-57, 164] and scales to [0, 1].\n","1. `CropForegroundd` removes all zero borders to focus on the valid body area of the images and labels.\n","1. `RandCropByPosNegLabeld` randomly crop patch samples from big image based on pos / neg ratio.  \n","The image centers of negative samples must be in valid body area.\n","1. `RandAffined` efficiently performs `rotate`, `scale`, `shear`, `translate`, etc. together based on PyTorch affine transform."]},{"cell_type":"code","execution_count":null,"metadata":{"id":"Jz3EcyN5XB3r"},"outputs":[],"source":["train_transforms = Compose(\n","    [\n","        LoadImaged(keys=[\"image\", \"label\"]),\n","        EnsureChannelFirstd(keys=[\"image\", \"label\"]),\n","        ScaleIntensityRanged(\n","            keys=[\"image\"],\n","            a_min=-57,\n","            a_max=164,\n","            b_min=0.0,\n","            b_max=1.0,\n","            clip=True,\n","        ),\n","        CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n","        Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n","        Spacingd(keys=[\"image\", \"label\"], pixdim=(1.5, 1.5, 2.0), mode=(\"bilinear\", \"nearest\")),\n","        RandCropByPosNegLabeld(\n","            keys=[\"image\", \"label\"],\n","            label_key=\"label\",\n","            spatial_size=(96, 96, 96),\n","            pos=1,\n","            neg=1,\n","            num_samples=4,\n","            image_key=\"image\",\n","            image_threshold=0,\n","        )\n","    ]\n",")\n","val_transforms = Compose(\n","    [\n","        LoadImaged(keys=[\"image\", \"label\"]),\n","        EnsureChannelFirstd(keys=[\"image\", \"label\"]),\n","        ScaleIntensityRanged(\n","            keys=[\"image\"],\n","            a_min=-57,\n","            a_max=164,\n","            b_min=0.0,\n","            b_max=1.0,\n","            clip=True,\n","        ),\n","        CropForegroundd(keys=[\"image\", \"label\"], source_key=\"image\"),\n","        Orientationd(keys=[\"image\", \"label\"], axcodes=\"RAS\"),\n","        Spacingd(keys=[\"image\", \"label\"], pixdim=(1.5, 1.5, 2.0), mode=(\"bilinear\", \"nearest\")),\n","    ]\n",")"]},{"cell_type":"markdown","metadata":{"id":"PweYNXRzXB3r"},"source":["## Check setup in DataLoader"]},{"cell_type":"code","execution_count":null,"metadata":{"tags":[],"id":"bhCpJM6GXB3r"},"outputs":[],"source":["check_ds = Dataset(data=val_files, transform=val_transforms)\n","check_loader = DataLoader(check_ds, batch_size=1)\n","check_data = first(check_loader)\n","image, label = (check_data[\"image\"][0][0], check_data[\"label\"][0][0])\n","print(f\"image shape: {image.shape}, label shape: {label.shape}\")\n","# plot the slice [:, :, 80]\n","plt.figure(\"check\", (12, 6))\n","plt.subplot(1, 2, 1)\n","plt.title(\"image\")\n","plt.imshow(image[:, :, 80], cmap=\"gray\")\n","plt.subplot(1, 2, 2)\n","plt.title(\"label\")\n","plt.imshow(label[:, :, 80])\n","plt.show()"]},{"cell_type":"markdown","metadata":{"id":"9ImZ25yvXB3r"},"source":["## CacheDataset\n","\n","Here we use CacheDataset to accelerate training and validation process, it's 10x faster than the regular Dataset.  \n","To achieve best performance, set `cache_rate=1.0` to cache all the data, if memory is not enough, set lower value.  \n","Users can also set `cache_num` instead of `cache_rate`, will use the minimum value of the 2 settings.  \n","And set `num_workers` to enable multi-threads during caching.  \n","If want to to try the regular Dataset, just change to use the commented code below."]},{"cell_type":"code","execution_count":null,"metadata":{"tags":[],"id":"cvCZklzfXB3r"},"outputs":[],"source":["train_ds = CacheDataset(data=train_files, transform=train_transforms, cache_rate=1.0, num_workers=4)\n","train_loader = DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=4)\n","\n","val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_rate=1.0, num_workers=4)\n","val_loader = DataLoader(val_ds, batch_size=1, num_workers=4)"]},{"cell_type":"markdown","source":["# Loss Function Timings\n","Write custom losses: [credit](https://github.com/JunMa11/SegWithDistMap/blob/master/code/train_LA_HD.py)"],"metadata":{"id":"dqkyQ7n8CmiH"}},{"cell_type":"code","source":["data_path_dir = f\"/content/drive/MyDrive/dev/MolloiLab/distance-transforms-paper/data\""],"metadata":{"id":"dkpZX3zoyJOe"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["np_pred = np.random.choice([0, 1], size=(1, 1, 224, 224, 112)).astype(np.float32)\n","np_label = np.random.choice([0, 1], size=(1, 1, 224, 224, 112)).astype(np.float32)\n","\n","torch_pred = torch.tensor(np_pred).cuda()\n","torch_label = torch.tensor(np_label).cuda()"],"metadata":{"id":"eSBLcFy6FCF-"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Dice Loss"],"metadata":{"id":"NdClTrHTCru6"}},{"cell_type":"code","source":["dice_loss = DiceLoss(to_onehot_y=False, softmax=False)"],"metadata":{"id":"ucEwd4R5Arap"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["%%timeit\n","dice_loss(torch_pred, torch_label)"],"metadata":{"id":"gtSWHz0DFFAT"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["import statistics"],"metadata":{"id":"d6sEE3V4n79U"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Number of runs\n","number = 5\n","repeat = 10"],"metadata":{"id":"iy2gVq9GrpUU"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Use timeit.repeat to get multiple timing results\n","times_dice = timeit.repeat(lambda: dice_loss(torch_pred, torch_label), number=7, repeat=number)"],"metadata":{"id":"EYME5DSwruG-"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Calculate statistics\n","min_dice, std_dice = min(times_dice), statistics.stdev(times_dice)\n","min_dice, std_dice # seconds"],"metadata":{"id":"nvGDIf93n4OO"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Scipy HD Loss"],"metadata":{"id":"zL8REo_WCxM8"}},{"cell_type":"code","source":["def compute_dtm(img_gt, out_shape):\n","    \"\"\"\n","    compute the distance transform map of foreground in binary mask\n","    input: segmentation, shape = (batch_size, x, y, z)\n","    output: the foreground Distance Map (SDM)\n","    dtm(x) = 0; x in segmentation boundary\n","             inf|x-y|; x in segmentation\n","    \"\"\"\n","\n","    fg_dtm = np.zeros(out_shape)\n","\n","    for b in range(out_shape[0]): # batch size\n","        for c in range(1, out_shape[1]):\n","            posmask = img_gt[b].astype(bool)\n","            if posmask.any():\n","                posdis = distance_transform_edt(posmask)\n","                fg_dtm[b][c] = posdis\n","\n","    return fg_dtm"],"metadata":{"id":"PoUC__jMe15X"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["def hd_loss(seg_soft, gt, seg_dtm, gt_dtm):\n","    \"\"\"\n","    compute huasdorff distance loss for binary segmentation\n","    input: seg_soft: softmax results,  shape=(b,2,x,y,z)\n","           gt: ground truth, shape=(b,x,y,z)\n","           seg_dtm: segmentation distance transform map; shape=(b,2,x,y,z)\n","           gt_dtm: ground truth distance transform map; shape=(b,2,x,y,z)\n","    output: boundary_loss; sclar\n","    \"\"\"\n","\n","    delta_s = (seg_soft[:,1,...] - gt.float()) ** 2\n","    s_dtm = seg_dtm[:,1,...] ** 2\n","    g_dtm = gt_dtm[:,1,...] ** 2\n","    dtm = s_dtm + g_dtm\n","\n","    multipled = torch.einsum('bxyz, bxyz->bxyz', delta_s, dtm)\n","    hd_loss = multipled.mean()\n","\n","    return hd_loss"],"metadata":{"id":"KkpD0KxKewrN"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Convert torch_pred to shape (b, 2, x, y, z) for softmax output\n","torch_pred_soft = torch.stack((1 - torch_pred, torch_pred), dim=1).squeeze(2)"],"metadata":{"id":"5JpqR4m_JwHl"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["def time_hd_loss_scipy():\n","  gt_dtm_npy = compute_dtm(torch_label.cpu().numpy(), torch_pred_soft.shape)\n","  gt_dtm = torch.from_numpy(gt_dtm_npy).float().cuda(torch_pred_soft.device.index)\n","  seg_dtm_npy = compute_dtm(torch_pred_soft[:, 1, :, :, :].cpu().numpy() > 0.5, torch_pred_soft.shape)\n","  seg_dtm = torch.from_numpy(seg_dtm_npy).float().cuda(torch_pred_soft.device.index)\n","\n","  hd_loss(torch_pred_soft, torch_label[:, 0, :, :, :], seg_dtm, gt_dtm)"],"metadata":{"id":"mWt9aMssr8OR"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["times_hd_scipy = timeit.repeat(lambda: time_hd_loss_scipy(), number=number, repeat=repeat)"],"metadata":{"id":"-VgY0ZzDIZev"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Calculate statistics\n","min_hd_scipy, std_hd_scipy = min(times_hd_scipy), statistics.stdev(times_hd_scipy)\n","min_hd_scipy, std_hd_scipy # seconds"],"metadata":{"id":"8uVaRTohsNYn"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## PyDTs HD Loss (Proposed)"],"metadata":{"id":"LlD1gyP-Ctnv"}},{"cell_type":"code","source":["def compute_dtm_gpu(img_gt, out_shape):\n","    \"\"\"\n","    compute the distance transform map of foreground in binary mask\n","    input: segmentation, shape = (batch_size, x, y, z)\n","    output: the foreground Distance Map (SDM)\n","    dtm(x) = 0; x in segmentation boundary\n","             inf|x-y|; x in segmentation\n","    \"\"\"\n","\n","    # Convert img_gt to float if not already float\n","    if img_gt.dtype != torch.float32:\n","        img_gt = img_gt.float()\n","\n","    fg_dtm = torch.zeros(out_shape, dtype=torch.float32, device=img_gt.device)\n","\n","    for b in range(out_shape[0]):  # batch size\n","        for c in range(1, out_shape[1]):\n","            posmask = img_gt[b]\n","            if posmask.bool().any():\n","                posdis = transform_cuda(posmask)\n","                fg_dtm[b, c] = posdis\n","\n","    return fg_dtm.to(img_gt.dtype)"],"metadata":{"id":"RD3EJHpM6cJ7"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["def time_hd_loss_pydt():\n","  gt_dtm = compute_dtm_gpu(torch_label, torch_pred_soft.shape)\n","  seg_dtm = compute_dtm_gpu(torch_pred_soft[:, 1, :, :, :] > 0.5, torch_pred_soft.shape)\n","  hd_loss(torch_pred_soft, torch_label[:, 0, :, :, :], seg_dtm, gt_dtm)"],"metadata":{"id":"Yr1h4gX0syrl"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["times_hd_pydt = timeit.repeat(lambda: time_hd_loss_pydt(), number=number, repeat=repeat)"],"metadata":{"id":"zrm59X7JJ7w4"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["jl.seval(\"CUDA.GC.gc(true); CUDA.reclaim()\") # IMPORTANT, otherwise GPU RAM can overflow"],"metadata":{"id":"Tk7nhklIFF_R"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Calculate statistics\n","min_hd_pydt, std_hd_pydt = min(times_hd_pydt[1:-1]), statistics.stdev(times_hd_pydt[1:-1])\n","min_hd_pydt, std_hd_pydt # seconds"],"metadata":{"id":"QeW_BWjPs9fx"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Create a dictionary with the results\n","results = {\n","    'Loss Function': ['Dice Loss', 'HD Loss (SciPy)', 'HD Loss (PyDT)'],\n","    'Minimum Time (s)': [min_dice, min_hd_scipy, min_hd_pydt],\n","    'Standard Deviation (s)': [std_dice, std_hd_scipy, std_hd_pydt]\n","}\n","\n","# Create the DataFrame\n","df_pure_losses = pd.DataFrame(results)\n","\n","# Set 'Loss Function' as the index\n","df_pure_losses = df_pure_losses.set_index('Loss Function')\n","\n","# Format the numbers to 6 decimal places\n","# df_pure_losses = df_pure_losses.applymap(lambda x: f\"{x:.6f}\")\n","\n","# Display the DataFrame\n","print(df_pure_losses)"],"metadata":{"id":"QQANHx6ctcH0"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["# Save the DataFrame as a CSV file\n","df_pure_losses_path = f\"{data_path_dir}/hd_loss_pure_losses_timings.csv\"\n","df_pure_losses.to_csv(df_pure_losses_path)"],"metadata":{"id":"_VLO11NXt8-y"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","metadata":{"id":"y_Fb0uQ9XB3r"},"source":["# Training Loop Timings"]},{"cell_type":"code","source":["dice_loss = DiceLoss(to_onehot_y=True, softmax=True)"],"metadata":{"id":"bbje82tzuvdM"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Pure Dice Loss (Gold Standard)"],"metadata":{"id":"rR3hvTCU8jfY"}},{"cell_type":"code","execution_count":null,"metadata":{"scrolled":true,"tags":[],"id":"TNlMW9sQXB3r","collapsed":true},"outputs":[],"source":["# standard PyTorch program style: create UNet, DiceLoss and Adam optimizer\n","device = torch.device(\"cuda:0\")\n","model = UNet(\n","    spatial_dims=3,\n","    in_channels=1,\n","    out_channels=2,\n","    channels=(16, 32, 64, 128, 256),\n","    strides=(2, 2, 2, 2),\n","    num_res_units=2,\n","    norm=Norm.BATCH,\n",").to(device)\n","optimizer = torch.optim.Adam(model.parameters(), 1e-4)\n","\n","max_epochs = 10\n","model_data = []\n","epoch_loss_values = []\n","epoch_times = []\n","\n","for epoch in range(max_epochs):\n","    start_time = time.time()\n","    print(\"-\" * 10)\n","    print(f\"epoch {epoch + 1}/{max_epochs}\")\n","    model.train()\n","    epoch_loss = 0\n","    step = 0\n","    for batch_data in train_loader:\n","        step += 1\n","        inputs, labels = (\n","            batch_data[\"image\"].to(device),\n","            batch_data[\"label\"].to(device),\n","        )\n","        optimizer.zero_grad()\n","        outputs = model(inputs)\n","        loss = dice_loss(outputs, labels)\n","        loss.backward()\n","        optimizer.step()\n","        epoch_loss += loss.item()\n","    epoch_loss /= step\n","    epoch_loss_values.append(epoch_loss)\n","    print(f\"epoch {epoch + 1} average loss: {epoch_loss:.4f}\")\n","\n","    end_time = time.time()\n","    epoch_time = end_time - start_time\n","    epoch_times.append(epoch_time)\n","    print(f\"Epoch {epoch + 1} training time: {epoch_time:.2f} seconds\")\n","\n","# Print total training time, average time per epoch, and standard deviation\n","total_training_time = sum(epoch_times)\n","avg_epoch_time = np.mean(epoch_times)\n","std_epoch_time = np.std(epoch_times)\n","print(f\"\\nTotal training time: {total_training_time:.2f} seconds\")\n","print(f\"Average training time per epoch: {avg_epoch_time:.2f} seconds\")\n","print(f\"Standard deviation of training time per epoch: {std_epoch_time:.2f} seconds\")\n","\n","# Append the model's details and timings to the list\n","model_data.append({\n","    'Model': 'Plain Dice Loss',\n","    'Total Training Time (s)': total_training_time,\n","    'Avg Epoch Time (s)': avg_epoch_time,\n","    'Std Epoch Time (s)': std_epoch_time\n","})\n","\n","# Create a DataFrame from the model_data list\n","df_plain_dice = pd.DataFrame(model_data)\n","\n","# Save the DataFrame as a CSV file\n","df_dice_path = f\"{data_path_dir}/hd_loss_plain_dice_timing.csv\"\n","df_plain_dice.to_csv(df_dice_path)"]},{"cell_type":"code","source":["torch.cuda.empty_cache()"],"metadata":{"id":"0WVPz2SEvGAt"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## Scipy HD Loss + Dice Loss\n","\n","Run this for only 10 epochs, just to showcase the difference in training speed. Accuracy differences between `scipy.ndimage.distance_transform_edt` and `py_distance_transforms.transform_cuda` should be neglible"],"metadata":{"id":"CayyNWyv8Xnv"}},{"cell_type":"code","source":["# standard PyTorch program style: create UNet, DiceLoss and Adam optimizer\n","device = torch.device(\"cuda:0\")\n","model = UNet(\n","    spatial_dims=3,\n","    in_channels=1,\n","    out_channels=2,\n","    channels=(16, 32, 64, 128, 256),\n","    strides=(2, 2, 2, 2),\n","    num_res_units=2,\n","    norm=Norm.BATCH,\n",").to(device)\n","optimizer = torch.optim.Adam(model.parameters(), 1e-4)\n","\n","max_epochs = 10\n","model_data = []\n","epoch_loss_values = []\n","epoch_times = []\n","\n","alpha = 1.0\n","\n","for epoch in range(max_epochs):\n","    start_time = time.time()\n","    print(\"-\" * 10)\n","    print(f\"epoch {epoch + 1}/{max_epochs}\")\n","    model.train()\n","    epoch_loss = 0\n","    step = 0\n","    for batch_data in train_loader:\n","        step += 1\n","        inputs, labels = (\n","            batch_data[\"image\"].to(device),\n","            batch_data[\"label\"].to(device),\n","        )\n","        optimizer.zero_grad()\n","        outputs = model(inputs)\n","        loss_seg_dice = dice_loss(outputs, labels)\n","        outputs_soft = F.softmax(outputs, dim=1)\n","\n","        with torch.no_grad():\n","            gt_dtm_npy = compute_dtm(labels.cpu().numpy(), outputs_soft.shape)\n","            gt_dtm = torch.from_numpy(gt_dtm_npy).float().cuda(outputs_soft.device.index)\n","            seg_dtm_npy = compute_dtm(outputs_soft[:, 1, :, :, :].cpu().numpy()>0.5, outputs_soft.shape)\n","            seg_dtm = torch.from_numpy(seg_dtm_npy).float().cuda(outputs_soft.device.index)\n","\n","        loss_hd = hd_loss(outputs_soft, labels[:, 0, :, :, :], seg_dtm, gt_dtm)\n","        loss = alpha*loss_seg_dice + (1 - alpha) * loss_hd\n","        loss.backward()\n","        optimizer.step()\n","        epoch_loss += loss.item()\n","    epoch_loss /= step\n","    epoch_loss_values.append(epoch_loss)\n","    print(f\"epoch {epoch + 1} average loss: {epoch_loss:.4f}\")\n","\n","    end_time = time.time()\n","    epoch_time = end_time - start_time\n","    epoch_times.append(epoch_time)\n","    print(f\"Epoch {epoch + 1} training time: {epoch_time:.2f} seconds\")\n","\n","    alpha -= 0.001\n","    if alpha <= 0.001:\n","        alpha = 0.001\n","\n","# Print total training time, average time per epoch, and standard deviation\n","total_training_time = sum(epoch_times)\n","avg_epoch_time = np.mean(epoch_times)\n","std_epoch_time = np.std(epoch_times)\n","print(f\"\\nTotal training time: {total_training_time:.2f} seconds\")\n","print(f\"Average training time per epoch: {avg_epoch_time:.2f} seconds\")\n","print(f\"Standard deviation of training time per epoch: {std_epoch_time:.2f} seconds\")\n","\n","# Append the model's details and timings to the list\n","model_data.append({\n","    'Model': 'Scipy HD Loss + Dice Loss',\n","    'Total Training Time (s)': total_training_time,\n","    'Avg Epoch Time (s)': avg_epoch_time,\n","    'Std Epoch Time (s)': std_epoch_time\n","})\n","\n","# Create a DataFrame from the model_data list\n","df_hd_dice_scipy = pd.DataFrame(model_data)\n","\n","# Save the DataFrame as a CSV file\n","df_hd_dice_scipy_path = f\"{data_path_dir}/hd_loss_hd_dice_scipy_timing.csv\"\n","df_hd_dice_scipy.to_csv(df_hd_dice_scipy_path)"],"metadata":{"id":"D18ziIYi5dF4"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["torch.cuda.empty_cache()\n","torch.cuda.empty_cache()\n","torch.cuda.empty_cache()"],"metadata":{"id":"PGitf34xvV0W"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["## PyDTs HD Loss + Dice Loss (Proposed)"],"metadata":{"id":"FTizEIW48cpT"}},{"cell_type":"code","source":["# standard PyTorch program style: create UNet, DiceLoss and Adam optimizer\n","device = torch.device(\"cuda:0\")\n","model = UNet(\n","    spatial_dims=3,\n","    in_channels=1,\n","    out_channels=2,\n","    channels=(16, 32, 64, 128, 256),\n","    strides=(2, 2, 2, 2),\n","    num_res_units=2,\n","    norm=Norm.BATCH,\n",").to(device)\n","optimizer = torch.optim.Adam(model.parameters(), 1e-4)\n","\n","max_epochs = 10\n","model_data = []\n","epoch_loss_values = []\n","epoch_times = []\n","\n","alpha = 1.0\n","\n","for epoch in range(max_epochs):\n","    start_time = time.time()\n","    print(\"-\" * 10)\n","    print(f\"epoch {epoch + 1}/{max_epochs}\")\n","    model.train()\n","    epoch_loss = 0\n","    step = 0\n","    for batch_data in train_loader:\n","        step += 1\n","        inputs, labels = (\n","            batch_data[\"image\"].to(device),\n","            batch_data[\"label\"].to(device),\n","        )\n","        optimizer.zero_grad()\n","        outputs = model(inputs)\n","        loss_seg_dice = dice_loss(outputs, labels)\n","        outputs_soft = F.softmax(outputs, dim=1)\n","\n","        with torch.no_grad():\n","            gt_dtm = compute_dtm_gpu(labels, outputs_soft.shape)\n","            seg_dtm = compute_dtm_gpu(outputs_soft[:, 1, :, :, :]>0.5, outputs_soft.shape)\n","\n","        loss_hd = hd_loss(outputs_soft, labels[:, 0, :, :, :], seg_dtm, gt_dtm)\n","        loss = alpha*loss_seg_dice + (1 - alpha) * loss_hd\n","        loss.backward()\n","        optimizer.step()\n","        epoch_loss += loss.item()\n","    jl.seval(\"CUDA.GC.gc(true); CUDA.reclaim()\") # IMPORTANT, otherwise GPU RAM overflows, not a huge slowdown penalty either\n","    epoch_loss /= step\n","    epoch_loss_values.append(epoch_loss)\n","    print(f\"epoch {epoch + 1} average loss: {epoch_loss:.4f}\")\n","\n","    end_time = time.time()\n","    epoch_time = end_time - start_time\n","    epoch_times.append(epoch_time)\n","    print(f\"Epoch {epoch + 1} training time: {epoch_time:.2f} seconds\")\n","\n","    alpha -= 0.001\n","    if alpha <= 0.001:\n","        alpha = 0.001\n","\n","# Print total training time, average time per epoch, and standard deviation\n","total_training_time = sum(epoch_times)\n","avg_epoch_time = np.mean(epoch_times)\n","std_epoch_time = np.std(epoch_times)\n","print(f\"\\nTotal training time: {total_training_time:.2f} seconds\")\n","print(f\"Average training time per epoch: {avg_epoch_time:.2f} seconds\")\n","print(f\"Standard deviation of training time per epoch: {std_epoch_time:.2f} seconds\")\n","\n","# Append the model's details and timings to the list\n","model_data.append({\n","    'Model': 'PyDTs HD Loss + Dice Loss',\n","    'Total Training Time (s)': total_training_time,\n","    'Avg Epoch Time (s)': avg_epoch_time,\n","    'Std Epoch Time (s)': std_epoch_time\n","})\n","\n","# Create a DataFrame from the model_data list\n","df_hd_dice_pydt = pd.DataFrame(model_data)\n","\n","# Save the DataFrame as a CSV file\n","df_hd_dice_pydt_path = f\"{data_path_dir}/hd_loss_hd_dice_pydt_timing.csv\"\n","df_hd_dice_pydt.to_csv(df_hd_dice_pydt_path)"],"metadata":{"id":"UE5d0O0d6ln_"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":["torch.cuda.empty_cache()"],"metadata":{"id":"x5o4RqAtxNlt"},"execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.8.13"},"colab":{"provenance":[{"file_id":"1YNou2N6cywlosHSuBP1Yjj6RLUl-SfLV","timestamp":1719507672872},{"file_id":"https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/spleen_segmentation_3d.ipynb","timestamp":1716315236912}],"gpuType":"L4","machine_shape":"hm"},"accelerator":"GPU"},"nbformat":4,"nbformat_minor":0}