Repository: pytorch
Issue: Cannot get deterministic Mask RCNN without running out of CUDA memory
Labels: module: cuda,triaged,module: cublas,module: determinism

### 🐛 Describe the bug

I created a modified version of torchvision's pretrained Mask RCNN model to train on my own dataset, and I set the multiple random seeds to be able to compare different runs. However, I recently realized that the results of different runs using identical parameters have slightly different loss results, which compound over the epochs. 

I researched this issue and found (https://github.com/pytorch/pytorch/issues/15359). Following the information I got from that discussion, I updated my versions of pytorch and torchvision using ```pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121```.

When I first ran my code using 
```python
torch.use_deterministic_algorithms(True)
```
I got the following error:

```python
RuntimeError: Deterministic behavior was enabled with either `torch.use_deterministic_algorithms(True)` or 
`at::Context::setDeterministicAlgorithms(true)`, but this operation is not deterministic because it uses CuBLAS and you 
have CUDA >= 10.2. To enable deterministic behavior in this case, you must set an environment variable before running 
your PyTorch application: CUBLAS_WORKSPACE_CONFIG=:4096:8 or CUBLAS_WORKSPACE_CONFIG=:16:8. For more 
information, go to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
```

I expected that error message because of the new flag to detect non-deterministic algorithms. So, I changed my setup code to set the CUBLAS workspace this way:

```python
import random
import numpy as np
import os

os.environ["CUBLAS_WORKSPACE_CONFIG"]=":4096:8"
import torch

from skimage.io import imread  # this is from the scikit-learn library

from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection import MaskRCNN_ResNet50_FPN_Weights
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

from torch.utils.data import Dataset
from torch.utils.data import DataLoader
```

I then restarted my kernel and reran the model. The problem is that now, my model runs out of memory with the following error message:

```python
Traceback (most recent call last):

  File ~\.conda\envs\python_3_10_new\lib\site-packages\spyder_kernels\py3compat.py:356 in compat_exec
    exec(code, globals, locals)

  File c:\users\cgravelmiguel\documents\git_repo\gitlab\deep-learning-model-library\utils\temp_maskrcnn_test.py:210
    loss_dict = model(images, targets)

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torch\nn\modules\module.py:1511 in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torch\nn\modules\module.py:1520 in _call_impl
    return forward_call(*args, **kwargs)

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torchvision\models\detection\generalized_rcnn.py:105 in forward
    detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torch\nn\modules\module.py:1511 in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torch\nn\modules\module.py:1520 in _call_impl
    return forward_call(*args, **kwargs)

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torchvision\models\detection\roi_heads.py:817 in forward
    rcnn_loss_mask = maskrcnn_loss(mask_logits, mask_proposals, gt_masks, gt_labels, pos_matched_idxs)

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torchvision\models\detection\roi_heads.py:111 in maskrcnn_loss
    mask_targets = [

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torchvision\models\detection\roi_heads.py:112 in <listcomp>
    project_masks_on_boxes(m, p, i, discretization_size) for m, p, i in zip(gt_masks, proposals, mask_matched_idxs)

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torchvision\models\detection\roi_heads.py:94 in project_masks_on_boxes
    return roi_align(gt_masks, rois, (M, M), 1.0)[:, 0]

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torchvision\ops\roi_align.py:236 in roi_align
    return _roi_align(input, rois, spatial_scale, output_size[0], output_size[1], sampling_ratio, aligned)

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torchvision\ops\roi_align.py:168 in _roi_align
    val = _bilinear_interpolate(input, roi_batch_ind, y, x, ymask, xmask)  # [K, C, PH, PW, IY, IX]

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torchvision\ops\roi_align.py:62 in _bilinear_interpolate
    v1 = masked_index(y_low, x_low)

  File ~\.conda\envs\python_3_10_new\lib\site-packages\torchvision\ops\roi_align.py:55 in masked_index
    return input[

OutOfMemoryError: CUDA out of memory. Tried to allocate 74.77 GiB. GPU 0 has a total capacity of 15.99 GiB of which 
12.45 GiB is free. Of the allocated memory 1.21 GiB is allocated by PyTorch, and 76.11 MiB is reserved by PyTorch but 
unallocated. If reserved but unallocated memory is large try setting 
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory 
Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

```

This makes no sense because I'm running a test model with only 10 images loaded in batches of 2. Without CUBLAS, one epoch of this simplified model takes <1 second and I have run models with >2000 images without issues before I tried ```torch.use_deterministic_algorithms(True)```. Given the problem with CUBLAS, I wonder if this issue might be related to CUBLAS' loss of performance documented in https://github.com/pytorch/pytorch/issues/109856?

Now I wonder: am I doomed to use ```torch.use_deterministic_algorithms(False)``` for this model structure given that setting it to ```True``` without CUBLAS gives me the first error, and that using CUBLAS gives me the second error?

Here is my full code. I tried to simplify it as much as possible. I am also not advanced enough to create a fake dataset within the code, so I am attaching my training dataset to reproduce the error. Thank you.

```python
import random
import numpy as np
import os

os.environ["CUBLAS_WORKSPACE_CONFIG"]=":4096:8"
import torch

from skimage.io import imread  # this is from the scikit-learn library

from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection import MaskRCNN_ResNet50_FPN_Weights
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

from torch.utils.data import Dataset
from torch.utils.data import DataLoader

seed = 3
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.use_deterministic_algorithms(True)

class MaskBoxDataset(Dataset):
    '''
    A Pytorch Dataset class to load the images and compute the mask bounding boxes
    
    Returns
    ------------
    images: torch.Tensor of size (B, C, H, W)
    gt bboxes: torch.Tensor of size (B, max_objects, 4)
    gt classes: torch.Tensor of size (B, max_objects)
    '''
    def __init__(self,
                 inputs_dim1: list,
                 targets: list):
        self.inputs_dim1 = inputs_dim1
        self.targets = targets
        
    def __len__(self):
        return len(self.inputs_dim1)
    
    def __getitem__(self, index):
        
        # Select the sample
        input_dim1_ID = self.inputs_dim1[index]
        target_ID = self.targets[index]

        # Load input and target as numpy arrays
        image_dim1, mask = imread(input_dim1_ID), imread(target_ID)
               
        # Triplicate the images to create a 3D image
        image = np.stack((image_dim1, image_dim1, image_dim1), 2)
        
        # instances are encoded as different colors
        obj_ids = np.unique(mask)
        
        # first id is the background, so remove it
        obj_ids = obj_ids[1:]
        
        # I need to set this up to get the boxes
        masks = mask == obj_ids[:, None, None]
        
        # get bounding box coordinates for each mask
        num_objs = len(masks)
        bboxes = []
        for i in range(num_objs):   

            # identify the location of the object mask             
            pos = np.where(masks[i])
                
            # append the coordinates of the mask as bounding box coordinates
            xmin = np.min(pos[1])
            xmax = np.max(pos[1])
            ymin = np.min(pos[0])
            ymax = np.max(pos[0])
                    
            bboxes.append([xmin, ymin, xmax, ymax])
        
        # Transform the image to a float32 tensor with the dimensions in their proper place.
        image_transformed = torch.as_tensor(image, dtype = torch.float32)
        image_transformed = image_transformed.permute(2,0,1)
        
        # Format arrays for the dictionary required by Mask RCNN
        boxes = torch.as_tensor(bboxes, dtype=torch.float32)
        labels = torch.ones((num_objs,), dtype=torch.int64)
        masks = torch.as_tensor(masks, dtype=torch.uint8)
        image_id = torch.tensor([index])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])

        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area

        return image_transformed, target

# Load an instance segmentation model pre-trained on COCO and setthe number of its layers that will learn (2). The rest are frozen.
model = maskrcnn_resnet50_fpn(weights=MaskRCNN_ResNet50_FPN_Weights.DEFAULT,
                              trainable_backbone_layers = 2)

# Get the number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features

# Replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2)

# Now get the number of input features for the mask classifier
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256

# and replace the mask predictor with a new one
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
                                                   hidden_layer,
                                                   2)

# move model to device
model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))

train_dir_dim1 = 'For_troubleshooting/Input_SLRM10m_256'
mask_dir = 'For_troubleshooting/Target_10m_256'

# Get a list of filenames (tiles) that will be imported
filenames = [f for f in os.listdir(mask_dir) if f.endswith(".tif")]

# Add the paths to the filenames in each category
inputs_train_dim1 = [f'{train_dir_dim1}/{item}' for item in filenames] 
targets_train = [f'{mask_dir}/{item}' for item in filenames] 

def collate_fn(batch):
    return tuple(zip(*batch))

# Workflow to create the PyTorch dataset (these are still numpy arrays)
trn_ds = MaskBoxDataset(inputs_train_dim1, targets_train)

# Workflow to load the images in batches
trn_dl = DataLoader(trn_ds, batch_size=2, shuffle=True, collate_fn=collate_fn, num_workers = 0)

'''

#### CODE TO TRAIN model

'''

model.train()

# Identify the parameters that will learn
params = [p for p in model.parameters() if p.requires_grad]


# Define the optimizer
optimizer = torch.optim.Adam(params, lr=1e-3)

for epoch in range(2):
    print(f'EPOCH {epoch}')
    
    for batch, data in enumerate(trn_dl):
        
        # Update the learning parameters
        optimizer.zero_grad()
        
        # get the data in the batch and train the model on it
        images, targets = data
        
        # Lists the images in the batch and send them to the appropriate device
        images = list(image.to("cuda") for image in images)
        
        # Get a dictionary of all targets for each image and send them to the appropriate device
        targets = [{k: v.to("cuda") if isinstance(v, torch.Tensor) else v for k, v in t.items()} for t in targets]
    
        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values())
        
        # print losses to see if they are consistent between runs using the same parameters
        print(losses)
                    
        losses.backward()
        optimizer.step()

# THE END
```
[For_troubleshooting.zip](https://github.com/pytorch/pytorch/files/14349353/For_troubleshooting.zip)

### Versions

```python
Collecting environment information...
PyTorch version: 2.2.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Pro
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.13 | packaged by Anaconda, Inc. | (main, Sep 11 2023, 13:24:38) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.22631-SP0
Is CUDA available: True
CUDA runtime version: 12.1.105

CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA RTX A4000
Nvidia driver version: 537.70
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9

CurrentClockSpeed=3312

DeviceID=CPU0

Family=207

L2CacheSize=14336

L2CacheSpeed=

Manufacturer=GenuineIntel

MaxClockSpeed=3312

Name=Intel(R) Core(TM) i9-10940X CPU @ 3.30GHz

ProcessorType=3

Revision=21767

Versions of relevant libraries:
[pip3] flake8==6.0.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.3
[pip3] numpydoc==1.5.0
[pip3] torch==2.2.0+cu121
[pip3] torch-snippets==0.523
[pip3] torchaudio==2.2.0+cu121
[pip3] torchmetrics==1.3.1
[pip3] torchvision==0.17.0+cu121
[conda] numpy                     1.26.3                   pypi_0    pypi
[conda] numpydoc                  1.5.0           py310haa95532_0  
[conda] torch                     2.2.0+cu121              pypi_0    pypi
[conda] torch-snippets            0.523                    pypi_0    pypi
[conda] torchaudio                2.2.0+cu121              pypi_0    pypi
[conda] torchmetrics              1.3.1                    pypi_0    pypi
[conda] torchvision               0.17.0+cu121             pypi_0    pypi
```

cc @csarofeen @ptrblck @xwang233 @mruberry @kurtamohler