OWL-ViT memory usage grows linearly with each prediction
Author: zduey
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.21.1
- Platform: Linux-5.10.102.1-microsoft-standard-WSL2-x86_64-with-glibc2.29
- Python version: 3.8.11
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.12.1+cu102 (False)
- 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?: No
- Using distributed or parallel set-up in script?: No

### Who can help?

@sgugger @NielsRogge 

### Information

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

### Tasks

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

### Reproduction

```python
import torch

from torchvision.datasets import FakeData
from torchvision.transforms.functional import pil_to_tensor
from transformers import OwlViTProcessor, OwlViTForObjectDetection

text_prompts = ["a photo of a cat", "a photo of a dog"]
dataset = FakeData(size=50, image_size=(3, 28, 28), transform=pil_to_tensor)
processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
target_sizes = torch.Tensor([[28, 28]])

for image, _ in dataset:
    inputs = processor(text=[text_prompts], images=image, return_tensors="pt")
    outputs = model(**inputs)
    _ = processor.post_process(outputs=outputs, target_sizes=target_sizes)[0]
```

### Expected behavior

I expect to be able to generate predictions from the OwlViTForObjectDetection model in a loop without memory usage increasing by ~1GB on each call to the model (line 15). Below, I've included a plot of memory usage over time. I profiled the code using `memory_profiler` to determine that it is the call to the model (not the processing or post processing) that seems to be the culprit.

![memory_usage](https://user-images.githubusercontent.com/3272567/184638200-834094c0-3dad-4b51-b076-b9229ca33534.png)



Comments:
cc @alaradirik as well :)

Thank you @zduey, I'm looking into this! 

@alaradirik - You are likely further along isolating the problem than I am, so please ignore if this is a distraction. The issue seems specific to repeated calls to the `OwlViTForObjectDetection` forward pass. 

The following results in the same memory usage pattern as the initial snippet. It differs from the original mainly in that the processor is not called as part of the loop.

```python
import requests
from PIL import Image
from transformers import OwlViTForObjectDetection, OwlViTProcessor
from tqdm import trange

text_prompts = ["a photo of a cat", "a photo of a dog"]

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
inputs = processor(text=[text_prompts], images=image, return_tensors="pt")

for _ in trange(15):
    _ = model(**inputs)
```
![owl_vit_for_object_detection](https://user-images.githubusercontent.com/3272567/185410275-50480ab1-4d22-41fd-8db9-785a5cbfc4f1.png)

When I've stepped through the code more granularly, the jump in memory usage happens [here](https://github.com/huggingface/transformers/blob/v4.21.1/src/transformers/models/owlvit/modeling_owlvit.py#L1259), which is just a call to `OwlVitModel`.  However, calls to   `OwlVitModel`  directly seem to have constant memory usage:

```python
from PIL import Image
import requests
from transformers import OwlViTProcessor, OwlViTModel

model = OwlViTModel.from_pretrained("google/owlvit-base-patch32")
processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)
inputs = processor(
    text=[["a photo of a cat", "a photo of a dog"]], images=image, return_tensors="pt"
)

for _ in range(50):
    _ = model(**inputs)
```
![owl_vit_model](https://user-images.githubusercontent.com/3272567/185410349-9c7bcad0-c2fc-4cfe-b02b-e9cad8c09e74.png)

Since I've run that separately without the same memory increase, the issue seems to be with how it gets called from `OWLViTForObjectDetection`, but I haven't been able to track it down. My initial guess (based on [this issue](https://discuss.pytorch.org/t/why-my-memory-keeps-increasing-at-every-iteration/118137)) was to look for a place where a tensor (one that is still attached to the computation graph) is being placed  into a standard python list.


Hi @zduey, thank you for the detailed analysis!

I'm more or less at the same point as you are, I confirmed that `OwlViTModel` is not the cause of the memory leak. I tracked the leak down to the calls to `image_text_embedder` method of `OwlViTForObjectDetection` and I'm working on the fix. I'm aiming to open a PR to fix this by tomorrow.

Just an update on the issue - this should be fixed when this [PR](https://github.com/huggingface/transformers/pull/18734) is merged. The memory leak was due to `OwlViTForObjectDetection` model making calls to `OwlViTModel`s non-forward methods and hence keeping track of all computational graphs.

Here is the code I used to confirm the fix:
```
import requests
from PIL import Image
import torch
from transformers import OwlViTModel, OwlViTForObjectDetection, OwlViTProcessor

model = OwlViTForObjectDetection.from_pretrained("google/owlvit-base-patch32")
processor = OwlViTProcessor.from_pretrained("google/owlvit-base-patch32")

text_prompts = ["a photo of a cat", "a photo of a dog"]
url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)
inputs = processor(text=[text_prompts], images=image, return_tensors="pt")

for i in range(50):
    with torch.no_grad():
        _ = model(**inputs)
```
![Figure_1](https://user-images.githubusercontent.com/8944735/186166457-f09822cf-0853-44bf-9ef0-421823032d00.png)

I will close this issue once the PR is merged.


We all learn in the hard way, @alaradirik . Just a few weeks ago, @NielsRogge and me had the same issue 😆
Whenever there is PyTorch memory issue -> check `with torch.no_grad()` first.
 

@ydshieh so true! Took me a deep dive into PyTorch docs and a while to debug it

Thanks so much @alaradirik for getting together a fix for this so quickly!

Closing this issue as the fix PR is merged

I don't think this bug is fixed. 
I was trying to use the model without torch.no_grad and the memory still grows linearly with every prediction I make, so something is still leaking

Hi @AlonZolfi  You should use `with` `torch.no_grad`, not `without`

I understand this is a workaround. However, I need the gradients during inference for another calculation so I cannot use it with torch.no_grad

In this case, you will have to manage the gradients by yourself: when to save them, when to delete some of them, etc. It's out of the scope of the the library's GitHub issue pages however.