Added empty_cache function to clear memory for CPU, CUDA, and MPS devices; updated imports accordingly.

Author: ATATC

mipcandy/data/__init__.py

@@ -5,7 +5,7 @@
 from mipcandy.data.geometric import ensure_num_dimensions, orthographic_views, aggregate_orthographic_views, crop
 from mipcandy.data.inspection import InspectionAnnotation, InspectionAnnotations, load_inspection_annotations, \
     inspect, ROIDataset, RandomROIDataset
-from mipcandy.data.io import fast_save, fast_load, resample_to_isotropic, load_image, save_image
+from mipcandy.data.io import fast_save, fast_load, resample_to_isotropic, load_image, save_image, empty_cache
 from mipcandy.data.sliding_window import do_sliding_window, revert_sliding_window, slide_dataset, \
     UnsupervisedSWDataset, SupervisedSWDataset
 from mipcandy.data.transform import JointTransform, MONAITransform

mipcandy/data/io.py

@@ -1,3 +1,4 @@
+from gc import collect
 from math import floor
 from os import PathLike
 
@@ -56,3 +57,13 @@ def save_image(image: torch.Tensor, path: str | PathLike[str]) -> None:
         image = auto_convert(ensure_num_dimensions(image, 3)).to(torch.uint8).permute(1, 2, 0)
         return SpITK.WriteImage(SpITK.GetImageFromArray(image.detach().cpu().numpy(), isVector=True), path)
     raise NotImplementedError(f"Unsupported file type: {path}")
+
+
+def empty_cache(device: Device) -> None:
+    match torch.device(device).type:
+        case "cpu":
+            collect()
+        case "cuda":
+            torch.cuda.empty_cache()
+        case "mps":
+            torch.mps.empty_cache()