Potential memory leakage of TensorFlow Swin model on kaggle!
### System Info

Info:

```
Framework: TensorFlow 2 (Keras)
Version: 2.6
OS: Kaggle
``` 


### Who can help?

[Swin Model Card](https://huggingface.co/microsoft/swin-small-patch4-window7-224) @amyeroberts
TensorFlow: @Rocketknight1
Vision: @NielsRogge, @sgugger

### Information

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

### Tasks

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

### Reproduction

A recent [kaggle competition](https://www.kaggle.com/competitions/google-universal-image-embedding) (hosted by Google), I tried to use pretrained `tf` swin transformer model from hugging face but even with the base model, I consistently received out of memory error.  Below is the submission status with a `base_tf_swin` model.

![image](https://user-images.githubusercontent.com/17668390/181924503-159801b3-f74a-418b-983f-2f68d31f87b2.png)

Some note:

- Other framework like pytorch works fine here.
- Other than this model, much larger model like  `tf_convnext_xlarge` is able to run without OOM. 

So, I'm assuming there might be some potential memory leakage in `tf_swin` implementation. Below is the code I use to build the complete model.

```python
id = "microsoft/swin-base-patch4-window7-224-in22k"

from transformers import AutoFeatureExtractor, TFSwinModel
feature_extractor = AutoFeatureExtractor.from_pretrained(id)
```

```python
inputs = keras.Input(shape=(None, None, 3), dtype='uint8')
mode_inputs = tf.cast(inputs, tf.float32)

mode_inputs = keras.layers.Resizing(*INPUT_SHAPE)(mode_inputs)
mode_inputs = keras.layers.Rescaling(scale=1.0 / 255)(mode_inputs)
mode_inputs = keras.layers.Normalization(
    mean=feature_extractor.image_mean,
    variance=[x ** 2 for x in feature_extractor.image_std ],
    axis=3
)(mode_inputs)
mode_inputs = keras.layers.Permute(dims=(3, 1, 2))(mode_inputs)

tf_huggingface_module = TFSwinModel.from_pretrained(id)
tf_huggingface_module.trainable = False
```
```python
logits = tf_huggingface_module(mode_inputs)
adv_logits = keras.Dense(64)(logits.pooler_output)

outputs = keras.layers.Lambda(
    lambda x: tf.math.l2_normalize(x, axis=-1), name='embedding_norm'
)(adv_logits)

tf_huggingface_classifier = keras.Model(inputs, outputs)
```



### Expected behavior

It should work like other model. To reproduce the issue exactly, (in the worst case), you may need to run it on kaggle platform. Kaggle submission status (as shown in the above diagram) is not very descriptive other than just showing submission status :(. Mainly, I like to know what could be the cause of it and any possible solution.

Hi @innat, thanks for flagging this! 

In order to help figure out what's causing the problem and possible solutions, could you please answer the following questions: 

* Could you give the version of transformers you're using and any other relevant packages? 
* Does the notebook run successfully before entering it as a submission? If not, what line of code causes the failure? 
* Could you give details on the checkpoint used for convnext? Can you confirm the convnext model works with the exact same pipeline? 
* When you said other frameworks work fine - can you confirm that you were able to use the equivalent Swin PyTorch model on the same swin checkpoint? 

What would help most and answer all of these would be a saved kaggle notebook that you could share.<|||||>Can you help me build an app

On Mon, Aug 1, 2022, 6:47 PM amyeroberts ***@***.***> wrote:

> Hi @innat <https://github.com/innat>, thanks for flagging this!
>
> In order to help figure out what's causing the problem and possible
> solutions, could you please answer the following questions:
>
>    - Could you give the version of transformers you're using and any
>    other relevant packages?
>    - Does the notebook run successfully before entering it as a
>    submission? If not, what line of code causes the failure?
>    - Could you give details on the checkpoint used for convnext? Can you
>    confirm the convnext model works with the exact same pipeline?
>    - When you said other frameworks work fine - can you confirm that you
>    were able to use the equivalent Swin PyTorch model on the same swin
>    checkpoint?
>
> What would help most and answer all of these would be a saved kaggle
> notebook that you could share.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/18376#issuecomment-1201519353>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AXV67Y5E7M3352TID62TFHLVXAETFANCNFSM55DNTU2A>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>
<|||||>Hello @amyeroberts; thanks for checking. To answer all of your query, 

> Could you give the version of transformers you're using and any other relevant packages?
1. It can be done, we will share a notebook file. Shortly, 

```python
tf.__version__, tfa.__version__, transformers.__version__
('2.6.4', '0.14.0', '4.22.0.dev0')
```

> Does the notebook run successfully before entering it as a submission? If not, what line of code causes the failure?
2. I hardly used hugging face vision model. It's kind of my first look of these vision models for current [on-going kaggle competition](https://www.kaggle.com/competitions/google-universal-image-embedding). 


> Could you give details on the checkpoint used for convnext? Can you confirm the convnext model works with the exact same pipeline?
3. Regarding the convnext checkpoint, yes, I can give you the exact file and reproducible code. And I confirm that hugging face convnext (larger one) runs fine whereas tiny swin gives OOM.


> When you said other frameworks work fine - can you confirm that you were able to use the equivalent Swin PyTorch model on the same swin checkpoint?
4. I should have elaborate more. I'm not PyTorch 1st user. Swin PyTorch model works fine is reported by other practitioners. 

---

> What would help most and answer all of these would be a saved kaggle notebook that you could share.

[Notebook Files](https://gist.github.com/innat/edf5d2c64d55e341efaee2884a8536e8)

It contains TensorFlow ConvNeXt and Swin Model pipelines and relevant package's version. The modeling strategy, saving, and submission process is followed according to the rules. The [evaluation page](https://www.kaggle.com/competitions/google-universal-image-embedding/overview/evaluation) also describes how they evaluate both framework and expected modeling approach. Hope it helps.
<|||||>Thank you for your response.

On Tue, Aug 2, 2022, 12:53 PM Mohammed Innat ***@***.***>
wrote:

> Hello @amyeroberts <https://github.com/amyeroberts>; thanks for checking.
> To answer all of your query,
>
> Could you give the version of transformers you're using and any other
> relevant packages?
>
>
>    1. It can be done, we will share a notebook file. Shortly,
>
> tf.__version__, tfa.__version__, transformers.__version__
> ('2.6.4', '0.14.0', '4.22.0.dev0')
>
> Does the notebook run successfully before entering it as a submission? If
> not, what line of code causes the failure?
>
>
>    1. I hardly used hugging face vision model. It's kind of my first look
>    of these vision models for current on-going kaggle competition
>    <https://www.kaggle.com/competitions/google-universal-image-embedding>.
>
> Could you give details on the checkpoint used for convnext? Can you
> confirm the convnext model works with the exact same pipeline?
>
>
>    1. Regarding the convnext checkpoint, yes, I can give you the exact
>    file and reproducible code. And I confirm that hugging face convnext
>    (larger one) runs fine whereas tiny swin gives OOM.
>
> When you said other frameworks work fine - can you confirm that you were
> able to use the equivalent Swin PyTorch model on the same swin checkpoint?
>
>
>    1. I should have elaborate more. I'm not PyTorch 1st user. Swin
>    PyTorch model works fine is reported by other practitioners.
>
> ------------------------------
>
> What would help most and answer all of these would be a saved kaggle
> notebook that you could share.
>
> Notebook Files
> <https://gist.github.com/innat/edf5d2c64d55e341efaee2884a8536e8>
>
> It contains TensorFlow ConvNeXt and Swin Model pipelines and relevant
> package's version. The modeling strategy, saving, and submission process is
> followed according to the rules. The evaluation page
> <https://www.kaggle.com/competitions/google-universal-image-embedding/overview/evaluation>
> also describes how they evaluate both framework and expected modeling
> approach. Hope it helps.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/18376#issuecomment-1202385475>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AXV67Y2MCP2O5YEEYZPM7TDVXED2DANCNFSM55DNTU2A>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>
<|||||>Hi @innat, thank you for all of your detailed responses and for sharing the notebook. 

I ran the notebook in kaggle and was able to save out the model with the checkpoint you used in your first example: `"microsoft/swin-base-patch4-window7-224-in22k"`

The notebook is here: https://www.kaggle.com/code/aeroberts4444/test-swin-saving/notebook

Are you able to run the notebook you shared on kaggle? Or do you still hit the OOM?<|||||>@amyeroberts Thanks for running the code. 

Yes, if you run the code that I shared, you won't see any OOM effect instant. As I said, I tried to submit two model from hugging-face (`"microsoft/swin-tiny-patch4-window7-224"` and `"facebook/convnext-large-224-22k-1k"`) to [this](https://www.kaggle.com/competitions/google-universal-image-embedding/overview/evaluation) competition. 

The convnext is comparatively much larger than tiny swin, but in the inference time, the submission status always exceed the allowed compute resource for tiny swin but works fine for large convnext model. That's why I kind of have **weak assumption** that, there may be some issue with swin implementation.  Also, later I realized that pytorch practitioners use `timm` version of swin model, and not from `huggingface` and no issue found about OOM with that.   

This competition is unique (no training or test data is provided), so it might be hard to debug the root cause. Please let me know if its out of scope to address such issue. 
<|||||>Hi @innat, thanks for clarifying. It's certainly a problem if there's a memory leak and one we'd want to address. I'm going to continue to look into this. As you said, because of the nature of kaggle and the competition it can be hard to debug. As such, it might take some time before I manage to figure out if there's a problem, what it is and how to solve. <|||||>@amyeroberts Thanks for your cordial support. I also informed competition host (googler), [HERE](https://www.kaggle.com/competitions/google-universal-image-embedding/discussion/336534#1883421), but no response yet.

cc @kfrancischen <|||||>Hi @innat. As mentioned above it's quite hard to debug without know what's happening during submission and logs from the kaggle notebook. My current best guess is it's due to the size of the saved Swin model. 

Using your script to create and save out a model, I looked at the sizes across different checkpoints: 
```
"microsoft/resnet-50"                              # 23,561,152 params
"google/vit-base-patch16-224-in21k"                # 86,389,248 params
"microsoft/swin-base-patch4-window7-224-in22k"     # 86,743,224 params
"microsoft/swin-tiny-patch4-window7-224"           # 27,519,354 params
"facebook/convnext-large-224-22k-1k"               # 196,230,336 params
```

```
tf_hf_classifier_convnext_large_224_22k_1k:
total 25712
drwxr-xr-x   6 amyroberts  staff   192B 10 Aug 13:13 .
drwxr-xr-x  24 amyroberts  staff   768B 10 Aug 13:13 ..
drwxr-xr-x   2 amyroberts  staff    64B 10 Aug 13:13 assets
-rw-r--r--   1 amyroberts  staff   510K 10 Aug 13:13 keras_metadata.pb
-rw-r--r--   1 amyroberts  staff    12M 10 Aug 13:13 saved_model.pb
drwxr-xr-x   4 amyroberts  staff   128B 10 Aug 13:13 variables

tf_hf_classifier_resnet_50:
total 12048
drwxr-xr-x   6 amyroberts  staff   192B 10 Aug 12:51 .
drwxr-xr-x  24 amyroberts  staff   768B 10 Aug 13:13 ..
drwxr-xr-x   2 amyroberts  staff    64B 10 Aug 12:51 assets
-rw-r--r--   1 amyroberts  staff   488K 10 Aug 12:51 keras_metadata.pb
-rw-r--r--   1 amyroberts  staff   5.4M 10 Aug 12:51 saved_model.pb
drwxr-xr-x   4 amyroberts  staff   128B 10 Aug 12:51 variables

tf_hf_classifier_swin_base_patch4_window7_224_in22k:
total 179216
drwxr-xr-x   6 amyroberts  staff   192B 10 Aug 13:00 .
drwxr-xr-x  24 amyroberts  staff   768B 10 Aug 13:13 ..
drwxr-xr-x   2 amyroberts  staff    64B 10 Aug 12:59 assets
-rw-r--r--   1 amyroberts  staff   7.4M 10 Aug 13:00 keras_metadata.pb
-rw-r--r--   1 amyroberts  staff    80M 10 Aug 13:00 saved_model.pb
drwxr-xr-x   4 amyroberts  staff   128B 10 Aug 12:59 variables

tf_hf_classifier_swin_tiny_patch4_window7_224:
total 83944
drwxr-xr-x   6 amyroberts  staff   192B 10 Aug 13:09 .
drwxr-xr-x  24 amyroberts  staff   768B 10 Aug 13:13 ..
drwxr-xr-x   2 amyroberts  staff    64B 10 Aug 13:09 assets
-rw-r--r--   1 amyroberts  staff   474K 10 Aug 13:09 keras_metadata.pb
-rw-r--r--   1 amyroberts  staff    41M 10 Aug 13:09 saved_model.pb
drwxr-xr-x   4 amyroberts  staff   128B 10 Aug 13:09 variables

tf_hf_classifier_vit_base_patch16_224_in21k:
total 21328
drwxr-xr-x   6 amyroberts  staff   192B 10 Aug 12:53 .
drwxr-xr-x  24 amyroberts  staff   768B 10 Aug 13:13 ..
drwxr-xr-x   2 amyroberts  staff    64B 10 Aug 12:53 assets
-rw-r--r--   1 amyroberts  staff   162K 10 Aug 12:53 keras_metadata.pb
-rw-r--r--   1 amyroberts  staff    10M 10 Aug 12:53 saved_model.pb
drwxr-xr-x   4 amyroberts  staff   128B 10 Aug 12:53 variables
```

I haven't dug much into why the model is so much larger. A cursory glance at the model graphs didn't reveal anything particularly surprising. <|||||>Randomly jumping in this thread :-)

- Are you able to reproduce this issue in a machine with similar spec as Kaggle machines?
- One way to narrow down to the root cause is to gradually remove some parts of code 
- From the provided notebook, we can't have any conclusion on memory leak. Memory leak refers to the memory usage increase during a repetition of the same call to a particular code block.
- Suggestion: try to see if this issue occurs during model saving, or the memory usage increases during inference time.<|||||>@amyeroberts 
Thanks for checking. I'll quickly check the size of these models in torch version. 

@kfrancischen
Your feedback is really much appreciate here. ([more info](https://www.kaggle.com/competitions/google-universal-image-embedding/discussion/336534#1859160))<|||||>I would suggest debug this in a VM outside Kaggle though. I remembered there is limited GPU/TPU hours per week on Kaggle. Don't waste your quota :-)<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.