Repository: fepegar/torchio

Issue:
Memory leak with TorchIO 0.20.1
### Is there an existing issue for this?

- [X] I have searched the existing issues

### Bug summary

A memory leak issue is observed in TorchIO version 20.0.1 during prolonged training sessions. The memory usage gradually increases over time, leading to an eventual out-of-memory error. This issue is not present in version 0.20.0, where memory usage remains stable under similar conditions. 

The following graphs show the evolution of RAM utilization during a basic DL training session, using the provided example code, with TorchIO version 0.20.0: 
![torch20_0_0](https://github.com/user-attachments/assets/65093145-91d1-459b-9cd2-48613daf4bdb)

Here with TorchIO 0.20.1:

![torchio20_0_1](https://github.com/user-attachments/assets/2a1d11fd-2ea3-417f-b1ff-8f4132220603)

Do you have any tips for preventing this problem with Torch 0.20.1?

Thanks :)

### Code for reproduction

```python
import copy
from typing import Any
import torch
import torch.utils.data
import torchio as tio
import pytorch_lightning as pl
import monai
import psutil
import matplotlib.pyplot as plt


ram_usage = []
timestamps = []


def get_ram_usage():
    return psutil.virtual_memory().percent

# %% Lightning module
class BasicPLModule(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.net = monai.networks.nets.Unet(spatial_dims=3, in_channels=1, out_channels=1, channels=(8, 16, 32), strides=(2, 2)).float()

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=0.0001)
        return optimizer

    def training_step(self, batch, batch_idx):
        inputs = batch['t1'][tio.DATA].float()
        targets = batch['brain'][tio.DATA].float()
        preds = self.net(inputs)
        return torch.nn.functional.mse_loss(preds, targets)

    def on_train_epoch_end(self):
        ram_usage.append(get_ram_usage())
        timestamps.append(self.current_epoch)


# %% Main program
if __name__ == '__main__':

    subject = tio.datasets.Colin27()
    # convert to instance of new subject class
    subject = tio.Subject(
        t1=subject['t1'],
        head=subject['head'],
        brain=subject['brain'],
    )
    subjects = 10 * [subject]
    transforms = tio.Compose([
        tio.ZNormalization(masking_method = tio.ZNormalization.mean),
        tio.transforms.CropOrPad(target_shape=192)
    ])
    subject_dataset = tio.SubjectsDataset(subjects, transforms)
    training_loader = tio.SubjectsLoader(subject_dataset, batch_size=1)
    reg_net = BasicPLModule()

    trainer_args = {
        'max_epochs': 200,
    }

    trainer_reg = pl.Trainer(**trainer_args)
    trainer_reg.fit(reg_net, training_loader)

    # Plotting the data
    plt.plot(timestamps, ram_usage)
    plt.title('RAM Usage Over Time')
    plt.xlabel('Time')
    plt.ylabel('RAM Usage (%)')
    plt.grid(True)
    plt.show()
```


### Actual outcome

After a long period :
Process finished with exit code 137 (interrupted by signal 9: SIGKILL)

### Error messages

_No response_

### Expected outcome

It's an example of code where the expected outcome is irrelevant, but the RAM utilization should remain relatively constant.

### System info

```shell
Platform:   Windows-10-10.0.22631-SP0
TorchIO:    0.20.1
PyTorch:    2.4.1+cu121
SimpleITK:  2.4.0 (ITK 5.4)
NumPy:      2.0.2
Python:     3.9.13 (tags/v3.9.13:6de2ca5, May 17 2022, 16:36:42) [MSC v.1929 64 bit (AMD64)]
```


Discussion:
sh-shahrokhi: Hello, Is there any chance that this get merged?
nicoloesch: @sh-shahrokhi I am not entirely sure when this is getting merged into the main branch. In the meantime, you can still utilise my feature branch that alleviates the issue (also referenced in the PR). The only ``issue'' is that its 0.20.1 instead of the newest version of torchio
fepegar: Picking this up now.

Patch:
diff --git a/src/torchio/data/subject.py b/src/torchio/data/subject.py
index 75ade08..25f871c 100644
--- a/src/torchio/data/subject.py
+++ b/src/torchio/data/subject.py
@@ -5,7 +5,6 @@ import pprint
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 from typing import Any
-from typing import Callable
 from typing import Optional
 
 import numpy as np
@@ -69,9 +68,6 @@ class Subject(dict):
         )
         return string
 
-    def __copy__(self):
-        return _subject_copy_helper(self, type(self))
-
     def __len__(self):
         return len(self.get_images(intensity_only=False))
 
@@ -426,25 +422,3 @@ class Subject(dict):
         from ..visualization import plot_subject  # avoid circular import
 
         plot_subject(self, **kwargs)
-
-
-def _subject_copy_helper(
-    old_obj: Subject,
-    new_subj_cls: Callable[[dict[str, Any]], Subject],
-):
-    result_dict = {}
-    for key, value in old_obj.items():
-        if isinstance(value, Image):
-            value = copy.copy(value)
-        else:
-            value = copy.deepcopy(value)
-        result_dict[key] = value
-
-    new = new_subj_cls(**result_dict)  # type: ignore[call-arg]
-    new.applied_transforms = old_obj.applied_transforms[:]
-    return new
-
-
-class _RawSubjectCopySubject(Subject):
-    def __copy__(self):
-        return _subject_copy_helper(self, Subject)
diff --git a/src/torchio/datasets/fpg.py b/src/torchio/datasets/fpg.py
index 3b8449c..4c36421 100644
--- a/src/torchio/datasets/fpg.py
+++ b/src/torchio/datasets/fpg.py
@@ -4,12 +4,12 @@ from ..constants import DATA_REPO
 from ..data import LabelMap
 from ..data import ScalarImage
 from ..data.io import read_matrix
-from ..data.subject import _RawSubjectCopySubject
+from ..data.subject import Subject
 from ..download import download_url
 from ..utils import get_torchio_cache_dir
 
 
-class FPG(_RawSubjectCopySubject):
+class FPG(Subject):
     """3T :math:`T_1`-weighted brain MRI and corresponding parcellation.
 
     Args:
diff --git a/src/torchio/datasets/itk_snap/itk_snap.py b/src/torchio/datasets/itk_snap/itk_snap.py
index 4450378..2f31626 100644
--- a/src/torchio/datasets/itk_snap/itk_snap.py
+++ b/src/torchio/datasets/itk_snap/itk_snap.py
@@ -2,12 +2,12 @@ import urllib.parse
 
 from ...data import LabelMap
 from ...data import ScalarImage
-from ...data.subject import _RawSubjectCopySubject
+from ...data.subject import Subject
 from ...download import download_and_extract_archive
 from ...utils import get_torchio_cache_dir
 
 
-class SubjectITKSNAP(_RawSubjectCopySubject):
+class SubjectITKSNAP(Subject):
     """ITK-SNAP Image Data Downloads.
 
     See `the ITK-SNAP website`_ for more information.
diff --git a/src/torchio/datasets/mni/mni.py b/src/torchio/datasets/mni/mni.py
index 5424e00..6b68499 100644
--- a/src/torchio/datasets/mni/mni.py
+++ b/src/torchio/datasets/mni/mni.py
@@ -1,8 +1,8 @@
-from ...data.subject import _RawSubjectCopySubject
+from ...data.subject import Subject
 from ...utils import get_torchio_cache_dir
 
 
-class SubjectMNI(_RawSubjectCopySubject):
+class SubjectMNI(Subject):
     """Atlases from the Montreal Neurological Institute (MNI).
 
     See `the website <https://nist.mni.mcgill.ca/?page_id=714>`_ for more
diff --git a/src/torchio/datasets/slicer.py b/src/torchio/datasets/slicer.py
index a66cadc..c29f09a 100644
--- a/src/torchio/datasets/slicer.py
+++ b/src/torchio/datasets/slicer.py
@@ -1,7 +1,7 @@
 import urllib.parse
 
 from ..data import ScalarImage
-from ..data.subject import _RawSubjectCopySubject
+from ..data.subject import Subject
 from ..download import download_url
 from ..utils import get_torchio_cache_dir
 
@@ -36,7 +36,7 @@ URLS_DICT = {
 }
 
 
-class Slicer(_RawSubjectCopySubject):
+class Slicer(Subject):
     """Sample data provided by `3D Slicer <https://www.slicer.org/>`_.
 
     See `the Slicer wiki <https://www.slicer.org/wiki/SampleData>`_
diff --git a/src/torchio/transforms/augmentation/composition.py b/src/torchio/transforms/augmentation/composition.py
index 9a02a84..2e8b7b3 100644
--- a/src/torchio/transforms/augmentation/composition.py
+++ b/src/torchio/transforms/augmentation/composition.py
@@ -44,7 +44,7 @@ class Compose(Transform):
     def __repr__(self) -> str:
         return f'{self.name}({self.transforms})'
 
-    def get_base_args(self) -> Dict:
+    def get_base_args(self) -> dict:
         init_args = super().get_base_args()
         if 'parse_input' in init_args:
             init_args.pop('parse_input')
@@ -109,7 +109,7 @@ class OneOf(RandomTransform):
         super().__init__(parse_input=False, **kwargs)
         self.transforms_dict = self._get_transforms_dict(transforms)
 
-    def get_base_args(self) -> Dict:
+    def get_base_args(self) -> dict:
         init_args = super().get_base_args()
         if 'parse_input' in init_args:
             init_args.pop('parse_input')
diff --git a/src/torchio/transforms/transform.py b/src/torchio/transforms/transform.py
index e0c94d1..67b67cc 100644
--- a/src/torchio/transforms/transform.py
+++ b/src/torchio/transforms/transform.py
@@ -155,9 +155,9 @@ class Transform(ABC):
         if self.keep is not None:
             images_to_keep = {}
             for name, new_name in self.keep.items():
-                images_to_keep[new_name] = copy.copy(subject[name])
+                images_to_keep[new_name] = copy.deepcopy(subject[name])
         if self.copy:
-            subject = copy.copy(subject)
+            subject = copy.deepcopy(subject)
         with np.errstate(all='raise', under='ignore'):
             transformed = self.apply_transform(subject)
         if self.keep is not None:


Test patch:
diff --git a/tests/data/test_subject.py b/tests/data/test_subject.py
index 2062e4d..f831018 100644
--- a/tests/data/test_subject.py
+++ b/tests/data/test_subject.py
@@ -186,3 +186,29 @@ class TestSubject(TorchioTestCase):
         loader = tio.SubjectsLoader(subjects, batch_size=4)
         batch = next(iter(loader))
         assert batch.__class__ is dict
+
+    def test_deep_copy_subject(self):
+        sub_copy = copy.deepcopy(self.sample_subject)
+        assert isinstance(sub_copy, tio.data.Subject)
+
+        new_tensor = torch.ones_like(sub_copy['t1'].data)
+        sub_copy['t1'].set_data(new_tensor)
+        # The data of the original subject should not be modified
+        assert not torch.allclose(sub_copy['t1'].data, self.sample_subject['t1'].data)
+
+    def test_shallow_copy_subject(self):
+        # We are creating a deep copy of the original subject first to not modify the original subject
+        copy_original_subj = copy.deepcopy(self.sample_subject)
+        sub_copy = copy.copy(copy_original_subj)
+        assert isinstance(sub_copy, tio.data.Subject)
+
+        new_tensor = torch.ones_like(sub_copy['t1'].data)
+        sub_copy['t1'].set_data(new_tensor)
+
+        # The data of both copies needs to be the same as we are using a shallow copy
+        assert torch.allclose(sub_copy['t1'].data, copy_original_subj['t1'].data)
+        # The data of the original subject should not be modified
+        assert not torch.allclose(sub_copy['t1'].data, self.sample_subject['t1'].data)
+        assert not torch.allclose(
+            copy_original_subj['t1'].data, self.sample_subject['t1'].data
+        )