Segment anything: CUDA out of memory
Author: rb-synth
Pull request: no

### System Info

- `transformers` version: 4.35.0
- Platform: Linux-5.15.0-1043-aws-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.1.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

@amyeroberts

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

The example script given on the [transformers README](https://huggingface.co/docs/transformers/v4.35.0/en/model_doc/sam#transformers.SamModel) gives a CUDA out of memory error on my 24GB GPU.

How can I reduce error consumption? I'm surprised that I'm hitting this limit, I didn't hit the same issue with the version of SAM pip installed directly from [their github](https://github.com/facebookresearch/segment-anything).

### Expected behavior

Code does not run out of CUDA memory.

Comments:
Hi @rb-synth, thanks for raising this issue! 

Interesting - looking at the [checkpoint on the hub](https://huggingface.co/facebook/sam-vit-huge/tree/main), the weights for `sam-vit-huge` are 2.56 GB so should fit into 24GB memory. 

Running the script on an NVIDIA A10G, I'm able to load the model onto the GPU and then similarly hit OOM on the forward pass of the model, so it's likely something in the logic of the model architecture cc @ArthurZucker @younesbelkada 


Hi! this is expected.
In the official repo they force-use `torch.no_grad()` in the forward pass: https://github.com/facebookresearch/segment-anything/blob/main/segment_anything/modeling/sam.py#L53 which we want avoid since we also support SAM fine-tuning (a fine-tuning notebook is in the listed resources on the docs).
I was able to run inference on a simple NVIDIA T4 16GB (free tier Google Colab) with the model loaded in full precision. So alternatively for faster inference you could run it in float16. All you need to do to avoid OOM is to wrap the forward pass with `torch.no_grad()` context manager or `torch.inference_mode()`

```diff
import torch
from PIL import Image
import requests
from transformers import SamModel, SamProcessor

device = "cuda" if torch.cuda.is_available() else "cpu"
model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")

img_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
input_points = [[[450, 600]]]  # 2D location of a window in the image

inputs = processor(raw_image, input_points=input_points, return_tensors="pt").to(device)

- outputs = model(**inputs)
+ with torch.no_grad():
+     outputs = model(**inputs)

masks = processor.image_processor.post_process_masks(
    outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
)
scores = outputs.iou_scores
```

If you want to run it in half-precision:

```diff
import torch
from PIL import Image
import requests
from transformers import SamModel, SamProcessor

device = "cuda" if torch.cuda.is_available() else "cpu"
- model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
+ model = SamModel.from_pretrained("facebook/sam-vit-huge", low_cpu_mem_usage=True, torch_dtype=torch.float16).to(device)
processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")

img_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")
input_points = [[[450, 600]]]  # 2D location of a window in the image

- inputs = processor(raw_image, input_points=input_points, return_tensors="pt").to(device)
+ inputs = processor(raw_image, input_points=input_points, return_tensors="pt").to(device, torch.float16)

- outputs = model(**inputs)
+ with torch.no_grad():
+    outputs = model(**inputs)

masks = processor.image_processor.post_process_masks(
-     outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
+    outputs.pred_masks.float().cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
)
scores = outputs.iou_scores
```

Updating in https://github.com/huggingface/transformers/pull/27268 the documentation to reflect this change. Closing this issue as it is expected and there is no bug in the modeling file, feel free to re-open if you have more questions! 

Thanks @younesbelkada for the detailed explanation and diff code snippet! 🔥