The results of run_mae.py pre-training were abnormal
Author: hhllxx1121
Pull request: no
Labels: name: Good Second Issue

### System Info

- `transformers` version: 4.39.0.dev0
- Platform: Linux-5.10.0-60.18.0.50.oe2203.x86_64-x86_64-with-glibc2.34
- Python version: 3.9.18
- Huggingface_hub version: 0.21.3
- Safetensors version: 0.4.2
- Accelerate version: 0.27.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.0.1+cu117 (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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@NielsRogge.

### Information

- [X] The official example scripts
- [ ] 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

1. ViT_MAE_visualization_demo: https://github.com/NielsRogge/Transformers-Tutorials/blob/master/ViTMAE/ViT_MAE_visualization_demo.ipynb
2. run_mae.py: https://github.com/huggingface/transformers/tree/main/examples/pytorch/image-pretraining

This is the result of loading the pre-trained model `facebook/vit-mae-huge` directly from huggingface.
![image](https://github.com/huggingface/transformers/assets/96508996/7ef67f11-c796-42ec-97a0-baa9101ae2fb)

This is the result of my continued pre_training using `facebook/vit-mae-huge`
![image](https://github.com/huggingface/transformers/assets/96508996/8e89d48c-2eb8-4bbd-8c15-dec6875f22e3)


I don't know why it's getting worse，The loss is decreasing from loss=1.7 at the beginning to loss=0.13. 

### Expected behavior

a

Comments:
I'll mark this issue as a "good second issue" for someone who could look into the details of the run_mae.py script.

For now I recommend to leverage the original MAE pre-training script: https://github.com/facebookresearch/mae/blob/main/main_pretrain.py. Things like learning rate, number of epochs, learning rate scheduler, etc. would need to be double-checked to make sure both match completely. It would also be interesting to check whether the same occurs with the original pre-training script.

The run_mae.py script was mainly created for demo purposes but has not yet been battle-tested in terms of a large-scale pre-training run.

Hello!

I've encountered the same unexpected behavior while training a model focused on facial recognition.
I started with a pretrained model ``facebook/vit-mae-base``, aiming to fine-tune it on a specific task related to faces.
I'm using a custom dataset from [``CristoJV/image-folder``](https://huggingface.co/datasets/CristoJV/image-folder)  as the one provided in the README ``nateraw/imagefolder``  is not compatible with the current version of ``datasets``.

This is the result of loading the pre-trained model facebook/vit-mae-base directly from huggingface.
![output](https://github.com/huggingface/transformers/assets/44086220/1e575c16-e19f-4344-92de-9f8e56430da1)
This is the result of my continued pre_training using facebook/vit-mae-base
![output2](https://github.com/huggingface/transformers/assets/44086220/01bbd94b-97e2-43ea-938e-5ae162381120)

How did you managed to train your ViTMAE? @hhllxx1121 

@NielsRogge, is there a direct way (implemented already) to convert the ``mae_pretrain_vit_base.pth`` model checkpoint for use with ``VitMAE (facebook/vit-mae-base)``? I'd like to utilize the ``main_pretrain.py`` script as suggested but I'm concerned that line 156 in main_pretraining.py may not work as expected when using the ``ViTMAE`` class. The code in question is:
```python
    model = models_mae.__dict__[args.model](norm_pix_loss=args.norm_pix_loss)
```
Thank you!


@CristoJV you can follow this guide: https://github.com/facebookresearch/mae/blob/main/PRETRAIN.md. The script supports a `--resume` flag, hence you can run the script with the following flag in order to further pre-train the model:

```
--resume "https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth"
```
This will automatically instantiate the weights of `facebook/vit-mae-base` into the script (in the original structure of the repository).

After pre-training, you can convert the original model to the HF format using the conversion script: https://github.com/huggingface/transformers/blob/main/src/transformers/models/vit_mae/convert_vit_mae_to_pytorch.py.



Yep 😄 I was looking for the convertion tool to speed up! 
Thank you!

Hi @NielsRogge and @hhllxx1121,

I have successfully managed to train the model using the original main_pretrain.py provided by the authors. Initially, I was getting the same results: the loss started higher than expected and suddenly decreased to low values, where it remained flat. This was observed after testing.

![output](https://github.com/huggingface/transformers/assets/44086220/2434a539-30d5-4a76-be5a-e782e319155c)

Then, I realized that I had enabled the pixel loss norm. After disabling it, I reran the training. This led me to wonder if the authors trained with the pixel loss norm disabled and, hence, it must be disabled when pretraining on that version. I have not yet tried to test this with the code developed by you, @NielsRogge, but I expect to see the same results.

Here is the image after retraining with the pixel norm loss disabled:
![output1](https://github.com/huggingface/transformers/assets/44086220/4b53dd3a-324d-44eb-a556-f6430f524ed6)

I'll provide an update after testing with the ``run_mae.py`` code. 😉

Another thing I forgot to mention: For pretraining, use this link:

"https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth"

instead of "https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth" which contains the decoder weights
 😉

> ```diff
>- --resume "https://dl.fbaipublicfiles.com/mae/pretrain/mae_pretrain_vit_base.pth"
> + --resume "https://dl.fbaipublicfiles.com/mae/visualize/mae_visualize_vit_base.pth"
> ```