Repository: deepchecks/deepchecks

Issue:
[BUG] [Vision] running ImageSegmentPerformance raises out of memory in both gpu and cpu devices
**Describe the bug**
`RuntimeError: CUDA out of memory. Tried to allocate 196.00 MiB (GPU 0; 15.90 GiB total capacity; 14.64 GiB already allocated; 145.75 MiB free; 14.82 GiB reserved in total by PyTorch)`


**To Reproduce**
Steps to reproduce the behavior:
1. Using birds 400 dataset from Kaggle  (Used 40 classes)
2. Run ImageSegmentPerformance
4. See error

**Expected behavior**
Possible to raise an error, but don't crash the notebook (try to catch OOM before hand)

**Screenshots**
If applicable, add screenshots to help explain your problem.

**Environment (please complete the following information):**
- OS: Kaggle Notebook
 - Python Version: 3.7.6
 - Deepchecks Version: 6.0



Discussion:
Can you attach the stacktrace? @benisraeldan 
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_34/533618478.py in <module>
      1 from deepchecks.vision.checks import ImageSegmentPerformance
      2 
----> 3 ImageSegmentPerformance().run(dataset=test_vd, model=model, device=device)

~/.local/lib/python3.7/site-packages/deepchecks/vision/base_checks.py in run(self, dataset, model, device, random_state)
     65             batch = Batch(batch, context, DatasetKind.TRAIN, batch_start_index)
     66             context.train.update_cache(batch)
---> 67             self.update(context, batch, DatasetKind.TRAIN)
     68             batch_start_index += len(batch)
     69 

~/.local/lib/python3.7/site-packages/deepchecks/vision/checks/performance/image_segment_performance.py in update(self, context, batch, dataset_kind)
     82         dataset = context.get_data_by_kind(dataset_kind)
     83         images = batch.images
---> 84         predictions = batch.predictions
     85         labels = batch.labels
     86 

~/.local/lib/python3.7/site-packages/deepchecks/vision/batch_wrapper.py in predictions(self)
     55             model = self._context.model
     56             self._context.assert_predictions_valid(self._dataset_kind)
---> 57             self._predictions = dataset.infer_on_batch(self._batch, model, self._context.device)
     58         return self._predictions
     59 

/tmp/ipykernel_34/670591851.py in infer_on_batch(self, batch, model, device)
     64 
     65         # Assuming the model returns the logits of the predictions, we need to convert them to probabilities.
---> 66         logits = model.to(device)(batch[0].to(device))
     67         return F.softmax(logits, dim=1)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/opt/conda/lib/python3.7/site-packages/torchvision/models/vgg.py in forward(self, x)
     47 
     48     def forward(self, x: torch.Tensor) -> torch.Tensor:
---> 49         x = self.features(x)
     50         x = self.avgpool(x)
     51         x = torch.flatten(x, 1)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/container.py in forward(self, input)
    137     def forward(self, input):
    138         for module in self:
--> 139             input = module(input)
    140         return input
    141 

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/conv.py in forward(self, input)
    441 
    442     def forward(self, input: Tensor) -> Tensor:
--> 443         return self._conv_forward(input, self.weight, self.bias)
    444 
    445 class Conv3d(_ConvNd):

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight, bias)
    438                             _pair(0), self.dilation, self.groups)
    439         return F.conv2d(input, weight, bias, self.stride,
--> 440                         self.padding, self.dilation, self.groups)
    441 
    442     def forward(self, input: Tensor) -> Tensor:

RuntimeError: CUDA out of memory. Tried to allocate 196.00 MiB (GPU 0; 15.90 GiB total capacity; 14.64 GiB already allocated; 145.75 MiB free; 14.82 GiB reserved in total by PyTorch)

```
How can I reproduce it? Can you please specify which dataset you tried to work on? and which model?
Maybe a kaggle link?
After investigation I think we should implement the following logic:
when `number_of_samples_to_infer_bins` is None, instead of caching a default of 1000 samples, we will cache X samples that will not use more than 10% of the available memory. We can measure the available memory using the `psutil` package.

Moving this issue to the 0.7 milestone.
Additional partial solution - move the cached batches to cpu always.

Resolving patch:
diff --git a/deepchecks/vision/checks/performance/image_segment_performance.py b/deepchecks/vision/checks/performance/image_segment_performance.py
--- a/deepchecks/vision/checks/performance/image_segment_performance.py
+++ b/deepchecks/vision/checks/performance/image_segment_performance.py
@@ -83,8 +83,8 @@ def initialize_run(self, context: Context, dataset_kind: DatasetKind):
     def update(self, context: Context, batch: Batch, dataset_kind: DatasetKind):
         """Update the bins by the image properties."""
         images = batch.images
-        predictions = batch.predictions
-        labels = batch.labels
+        predictions = [tens.detach() for tens in batch.predictions]
+        labels = [tens.detach() for tens in batch.labels]
 
         samples_for_bin: t.List = self._state['samples_for_binning']
         bins = self._state['bins']


Test patch: