Add FlaxVisionEncoderDecoderModel
Author: ydshieh
Pull request: yes

# What does this PR do?

Add `FlaxVisionEncoderDecoderModel`.

- Unlike the test for `FlaxVisionEncoderDecoderModel` where there is `test_bert2gpt2_summarization`, there is currently no official pretrained checkpoint for  `FlaxVisionEncoderDecoderModel`, and therefore there is no similar test for it. We can add this test once the PR is valid and used to train a `FlaxVisionEncoderDecoderModel`.

- Some tests involving the shapes of the `encoder_attentions` and `cross_attentions` require knowing the sequence lengths once the `pixel_values` is processed to a sequence. I don't find a built-in method in `ViT` providing this computation, so currently I just check the `batch_dim` and `hidden_dim`. Maybe it is good to think of this and add some utility functions to vision models. (This is computed as `self.num_patches` I think, but similar to the issue below, I don't have a clear idea how to access it for now).

- The test

https://github.com/huggingface/transformers/blob/35a3921180fe053ad474e137c0d738dfac543a8b/tests/test_modeling_flax_vision_encoder_decoder.py#L332
will fail since `FlaxVisionEncoderDecoderModel` doesn't have encoder/decoder attributes (I tried it locally by disable `@slow`). I could't find a way to access `FlaxVisionEncoderDecoderModel.module.encoder` and ``FlaxVisionEncoderDecoderModel.module.decoder`. The same issue exists for 

https://github.com/huggingface/transformers/blob/35a3921180fe053ad474e137c0d738dfac543a8b/tests/test_modeling_flax_encoder_decoder.py#L364
(It is a slow test, so won't show up when we push).

## Who can review?

@patil-suraj @patrickvonplaten 


Comments:
Hi @ydshieh,

Thanks for this PR, I will draft a `VisionEncoderDecoderModel` soon, as I would like to add it alongside the TrOCR model by Microsoft, which combines a vision encoder (BEiT) with an autoregressive text decoder (initialized from RoBERTa).

However, the current design of the `EncoderDecoderModels` does have a limitation: if the `hidden_size` of the encoder and decoder don't match, then they define a single projection layer to project the `encoder_hidden_states` to the same dimension as the decoder. However, in TrOCR, this is not how it's done: there, they project the `encoder_hidden_states` to the dimension of the decoder when defining the keys and values, in each layer separately.

Hence, I will set up a draft PR that allows both options (either a single projection layer, or projecting them in the keys/values).



> Hence, I will set up a draft PR that allows both options (either a single projection layer, or projecting them in the keys/values).

Hi, @NielsRogge . Maybe it's good to add this option to `VisionEncoderDecoderModel` as a first attempt, and apply to `EncoderDecoderModels` later? Nothing huge, but since you have a target `TrOCR model`, I think it's good to start with `VisionEncoderDecoderModel`.

Hi, can you perhaps also add support for `cross_attention_hidden_size` similar to the PyTorch implementation?

This is to support the case where, instead of just projecting the `encoder_hidden_states` to the same size as the decoder with a single projection layer, one projects them when creating keys and values in each cross-attention layer of the decoder.

We could perhaps, for consistency, add this also to the existing, text-only encoder-decoder model.

> 
> 
> Hi, can you perhaps also add support for `cross_attention_hidden_size` similar to the PyTorch implementation?
> 
> This is to support the case where, instead of just projecting the `encoder_hidden_states` to the same size as the decoder with a single projection layer, one projects them when creating keys and values in each cross-attention layer of the decoder.
> 
> We could perhaps, for consistency, add this also to the existing, text-only encoder-decoder model.

Sure, I will add it :)

For the text-only encoder-decoder, maybe in another PR though.

@NielsRogge , I add the projection layer to `FlaxVisionEncoderDecoderModel` when the condition meets. (and the PT/Flax equivalence test will check both cases)

It would be great if you can check the few minor changes I made in `modeling_vision_encoder_decoder.py` and `configuration_vision_encoder_decoder.py`.

@patrickvonplaten 
I added PT/Flax equivalence test + an image captioning pretrained model integration test: All passed.
(currently, I comment out some @slow tests to make sure all tests pass. I will clean up them later).

I also uploaded the PyTorch version of the new trained Flax's image captioning model, and added the same test for `VisionEncoderDecoderModel`.

https://github.com/huggingface/transformers/blob/05918efda1924510a4cf5b3ce0c7ca0b6bd8cf22/tests/test_modeling_flax_vision_encoder_decoder.py#L465

https://github.com/huggingface/transformers/blob/05918efda1924510a4cf5b3ce0c7ca0b6bd8cf22/tests/test_modeling_flax_vision_encoder_decoder.py#L332

I think this PR is ready for review.

@patil-suraj - feel free to merge once you're happy with the PR

Impressive work! I also saw your image captioning demo - so if I understand correctly, you first implemented a Vit2Gpt2 model yourself which you used for training (without using this new `FlaxVisionEncoderDecoderModel` class), and now you can load the weights of that model into it?

@NielsRogge I made a ViT+GPT2 during the Flax community week, and it had 2 major bugs, so gave really non-sense captions.
After that event, I decided to work on Encoder-Decoder models.

For the demo you saw, it does use the newly added `FlaxVisionEncoderDecoderModel` to load original pretrained ViT / GPT2 model, then finetuned on COCO 2017 :)

- added related lines for `AutoModelForVision2Seq`
- removed `test_configuration_tie`
- change `# @slow` to `@slow`

I removed `attention_mask` from `FlaxVisionEncoderDecoderModel.__call__`, and also from the inputs docstring of both `modeling_vision_encoder_decoder.py` and `modeling_flax_vision_encoder_decoder.py`.

cc @NielsRogge 

All reviews from @patil-suraj have been addressed, except for a few `raise NotImplementedError` in the test.

I rebased on master, and needed to fix a pt/flax equivalence test.

Some final clean-ups regarding the encoder attention mask and then we're good to go IMO :-)

> 
> 
> Some final clean-ups regarding the encoder attention mask and then we're good to go IMO :-)

Clean-ups done. I have checked again the slow tests and all good :-)
(other failed tests seem unrelated)

Awesome - thanks a lot @ydshieh! Could you maybe rebase (or merge master into your branch) to fix the failing tests? :-) There was a problem with TF tests on master

It's green now :-)

Hi, @patrickvonplaten @NielsRogge , about the blog post for image captioning mentioned in

https://github.com/huggingface/transformers/pull/14139#issuecomment-950868338

```
Also cc @ydshieh - If you're interested we could work on a nice blog post
(to be added here: https://huggingface.co/blog) on
how to leverage VisionEncoderDecoder for image captioning if you're interested :-)
```

I am thinking to read

https://huggingface.co/blog/encoder-decoder
https://huggingface.co/blog/fine-tune-wav2vec2-english

to get some idea about the writing.

My Flax/TPU training script however uses a custom `datasets` script to load a local COCO dataset, and there are something to modify to avoid the dataset processing speed issue.

I am just wondering if the blog post should include this training processes, or just contains how to use a pretrained model?

Pinged you on Slack, to not spam this PR :)

@patrickvonplaten I changed it back to `one` with copying you words as comment

https://github.com/huggingface/transformers/blob/d8aea41778e709ee6c6a9e8c84247d7cd9421c10/src/transformers/models/vision_encoder_decoder/modeling_flax_vision_encoder_decoder.py#L228-L233

The failure is unrelated to this PR, merging now.