Repository: mlcommons/GaNDLF

Issue:
Memory usage is exploding after adding subject consistency check
**Describe the bug**
When `torchio.Subject.check_consistent_*` checks are [called](https://github.com/CBICA/GaNDLF/blob/06e040627cb2ec1968ea3d81475ceeb0a7b44edc/GANDLF/data/ImagesFromDataFrame.py#L346), RAM usage explodes, most likely because of [this bug in PyTorch](https://github.com/pytorch/pytorch/issues/49394).

**To Reproduce**
Steps to reproduce the behavior:
1. Start training with a large number of subjects.
2. See memory usage linearly increasing and never going down

**Expected behavior**
Consistency check should happen once and then the image should be unloaded from memory.

**Screenshots**
If applicable, add screenshots to help explain your problem.

**GaNDLF Version**
0.0.10-dev

**Desktop (please complete the following information):**
N.A.

**Additional context**
N.A.



Discussion:

Resolving patch:
diff --git a/GANDLF/data/ImagesFromDataFrame.py b/GANDLF/data/ImagesFromDataFrame.py
--- a/GANDLF/data/ImagesFromDataFrame.py
+++ b/GANDLF/data/ImagesFromDataFrame.py
@@ -26,7 +26,7 @@
 from torchio import Image, Subject
 import SimpleITK as sitk
 
-# from GANDLF.utils import resize_image
+from GANDLF.utils import perform_sanity_check_on_subject
 from GANDLF.preprocessing import (
     NonZeroNormalizeOnMaskedRegion,
     CropExternalZeroplanes,
@@ -342,11 +342,13 @@ def ImagesFromDataFrame(dataframe, parameters, train):
             # Initializing the subject object using the dict
             subject = Subject(subject_dict)
             # https://github.com/fepegar/torchio/discussions/587#discussioncomment-928834
-            print("Checking consistency of images in subject '" + subject["subject_id"])
-            subject.check_consistent_affine()
-            subject.check_consistent_space()
-            subject.check_consistent_spatial_shape()
-            subject.check_consistent_orientation()
+            # this is causing memory usage to explode, see https://github.com/CBICA/GaNDLF/issues/128
+            print(
+                "Checking consistency of images in subject '"
+                + subject["subject_id"]
+                + "'"
+            )
+            perform_sanity_check_on_subject(subject, parameters)
 
             # # padding image, but only for label sampler, because we don't want to pad for uniform
             if "label" in sampler or "weight" in sampler:
diff --git a/GANDLF/utils.py b/GANDLF/utils.py
--- a/GANDLF/utils.py
+++ b/GANDLF/utils.py
@@ -578,3 +578,70 @@ def populate_channel_keys_in_params(data_loader, parameters):
         parameters["value_keys"] = value_keys
 
     return parameters
+
+
+def perform_sanity_check_on_subject(subject, parameters):
+    """
+    This function performs sanity check on the subject to ensure presence of consistent header information WITHOUT loading images into memory.
+
+    Args:
+        subject (torchio.Subject): The input subject.
+        parameters (dict): The parameters passed by the user yaml.
+
+    Returns:
+        bool: True if everything is okay.
+
+    Raises:
+        ValueError: Dimension mismatch in the images.
+        ValueError: Origin mismatch in the images.
+        ValueError: Orientation mismatch in the images.
+    """
+    # read the first image and save that for comparison
+    file_reader_base = None
+
+    import copy
+
+    list_for_comparison = copy.deepcopy(parameters["headers"]["channelHeaders"])
+    if "labelHeader" in parameters["headers"]:
+        list_for_comparison.append("label")
+
+    for key in list_for_comparison:
+        if file_reader_base is None:
+            file_reader_base = sitk.ImageFileReader()
+            file_reader_base.SetFileName(subject[str(key)]["path"])
+            file_reader_base.ReadImageInformation()
+        else:
+            # in this case, file_reader_base is ready
+            file_reader_current = sitk.ImageFileReader()
+            file_reader_current.SetFileName(subject[str(key)]["path"])
+            file_reader_current.ReadImageInformation()
+
+            if file_reader_base.GetDimension() != file_reader_current.GetDimension():
+                raise ValueError(
+                    "Dimensions for Subject '"
+                    + subject["subject_id"]
+                    + "' are not consistent."
+                )
+
+            if file_reader_base.GetOrigin() != file_reader_current.GetOrigin():
+                raise ValueError(
+                    "Origin for Subject '"
+                    + subject["subject_id"]
+                    + "' are not consistent."
+                )
+
+            if file_reader_base.GetDirection() != file_reader_current.GetDirection():
+                raise ValueError(
+                    "Orientation for Subject '"
+                    + subject["subject_id"]
+                    + "' are not consistent."
+                )
+
+            if file_reader_base.GetSpacing() != file_reader_current.GetSpacing():
+                raise ValueError(
+                    "Spacing for Subject '"
+                    + subject["subject_id"]
+                    + "' are not consistent."
+                )
+
+    return True


Test patch: