Add TFData2VecVision for semantic segmentation
Author: sayakpaul
Pull request: yes

This PR introduces `TFData2VecVisionForSemanticSegmentation` which takes the `TFData2VecVisionMainLayer` and appends the necessary layers for performing semantic segmentation along with loss computation (first one in this line?). 

**Notes**

* Thanks to @Rocketknight1 who implemented the adaptive average pooling layer. 
* Currently the model saving the tests (2 tests) are failing as soon as `TFData2VecVisionForSemanticSegmentation` class is introduced to `tests/models/test_modeling_tf_data2vec_vision.py`. Without that class, the test runs as expected. I would appreciate any help.
* As per discussed over Slack, [this class](https://github.com/huggingface/transformers/blob/main/src/transformers/models/data2vec/modeling_data2vec_vision.py#L882) should never have been subclassed from `nn.ModuleList`. It is currently leading a few idiosyncracies on the TF side (mainly related to naming of the layers). Once that is sorted out we can again revisit this `TFData2VecVisionForSemanticSegmentation` class and make the amends if needed. Happy to take the charge then. 
* I ran the tests locally with the following command: `RUN_SLOW=1 python -m pytest tests/models/data2vec/test_modeling_tf_data2vec_vision.py`. 

Here's the trace of the errors from running tests:

```
            model = model_class(config)
            model(self._prepare_for_class(inputs_dict, model_class))  # Model must be called before saving.
            # Let's load it from the disk to be sure we can use pretrained weights
            with tempfile.TemporaryDirectory() as tmpdirname:
>               model.save_pretrained(tmpdirname, saved_model=False)

tests/test_modeling_tf_common.py:693: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
src/transformers/modeling_tf_utils.py:1513: in save_pretrained
    self.save_weights(output_model_file)
../../.local/bin/.virtualenvs/hf/lib/python3.8/site-packages/keras/utils/traceback_utils.py:67: in error_handler
    raise e.with_traceback(filtered_tb) from None
../../.local/bin/.virtualenvs/hf/lib/python3.8/site-packages/h5py/_hl/group.py:149: in create_dataset
    dsid = dataset.make_new_dset(group, shape, dtype, data, name, **kwds)
../../.local/bin/.virtualenvs/hf/lib/python3.8/site-packages/h5py/_hl/dataset.py:142: in make_new_dset
    dset_id = h5d.create(parent.id, name, tid, sid, dcpl=dcpl)
h5py/_objects.pyx:54: in h5py._objects.with_phil.wrapper
    ???
h5py/_objects.pyx:55: in h5py._objects.with_phil.wrapper
    ???
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

>   ???
E   ValueError: Unable to create dataset (name already exists)

h5py/h5d.pyx:87: ValueError

...


           outputs = model(self._prepare_for_class(inputs_dict, model_class))
    
            with tempfile.TemporaryDirectory() as tmpdirname:
>               model.save_pretrained(tmpdirname, saved_model=False)

tests/test_modeling_tf_common.py:175: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
src/transformers/modeling_tf_utils.py:1513: in save_pretrained
    self.save_weights(output_model_file)
../../.local/bin/.virtualenvs/hf/lib/python3.8/site-packages/keras/utils/traceback_utils.py:67: in error_handler
    raise e.with_traceback(filtered_tb) from None
../../.local/bin/.virtualenvs/hf/lib/python3.8/site-packages/h5py/_hl/group.py:149: in create_dataset
    dsid = dataset.make_new_dset(group, shape, dtype, data, name, **kwds)
../../.local/bin/.virtualenvs/hf/lib/python3.8/site-packages/h5py/_hl/dataset.py:142: in make_new_dset
    dset_id = h5d.create(parent.id, name, tid, sid, dcpl=dcpl)
h5py/_objects.pyx:54: in h5py._objects.with_phil.wrapper
    ???
h5py/_objects.pyx:55: in h5py._objects.with_phil.wrapper
    ???
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

>   ???
E   ValueError: Unable to create dataset (name already exists)

h5py/h5d.pyx:87: ValueError
-------------------------------
```

Additionally, here's a little code for testing the segmentation class:

```py
from PIL import Image
import tensorflow as tf 

from src.transformers.models.data2vec.modeling_tf_data2vec_vision import (
    TFData2VecVisionForSemanticSegmentation
)
from transformers import BeitFeatureExtractor


def prepare_img():
    image = Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
    return image


feature_extractor = BeitFeatureExtractor.from_pretrained(
    "facebook/data2vec-vision-base-ft1k"
)
model = TFData2VecVisionForSemanticSegmentation.from_pretrained(
    "facebook/data2vec-vision-base",
)


image = prepare_img()
inputs = feature_extractor(images=image, return_tensors="tf")
batch_size, num_channels, height, width = inputs["pixel_values"].shape
inputs["labels"] = tf.zeros((batch_size, height, width))
outputs = model(**inputs)

print(outputs.logits.shape)
print(outputs.loss.shape)
```

@Rocketknight1 @sgugger 

Comments:
_The documentation is not available anymore as the PR was closed or merged._

> Thanks a lot for your PR! Note that on the pyramid pooling class, even if we change the PyTorch class to not subclass `ModuleList` anymore, it will still need to keep the same weight names, otherwise compatibility with any checkpoint on the Hub will be broken.

Absolutely. 

@Rocketknight1 a gentle ping 👀

Ah, I'm sorry! Will review it by tomorrow.

Hi, I just took a look over this! I suspect the issue with the tests is that there's something like a layer name collision when saving. In h5 files, weights are saved as 'datasets' , so this error is telling us that the weights are not uniquely named - the same 'dataset' name is being written to twice during saving, which means two layers share the same name.

Yes, I suspected something similar but couldn't figure out where the duplicate is coming from. Do you have any suggestions?

@Rocketknight1 

I suspect the issue is most likely related to the implementation of AdaptiveAvgPool I wrote - the practice of precomputing a constant sparse matrix like that is non-standard, and TF might be trying to save that Tensor somehow. Can you try replacing it with a 'dummy' layer that has the same output shape and seeing if the error goes away? If so, I can work on a different implementation for the layer - I have some ideas that I think will improve performance a lot, and they might also resolve the problem too.

> Can you try replacing it with a 'dummy' layer that has the same output shape and seeing if the error goes away? 

Sure. I will do it and get back. 

@Rocketknight1 this is what I did:

https://github.com/sayakpaul/transformers/blob/f9292cf2c47baf7eb264c98c6189ae503930130f/src/transformers/models/data2vec/modeling_tf_data2vec_vision.py#L1102

Same issue. 

@sayakpaul I used post-mortem debugging to isolate this - just add this to `TFData2VecVisionModelTest`:

```
  def test_save_load(self):
      try:
          super().test_save_load()
      except:
          import pdb
          pdb.post_mortem()
```

Then run the tests with `pytest --capture=no`. This will break into a debugger at the point of failure, and you can step up to the calling frame with `(u)p`.

From there, I can tell that the offending array has name `kernel:0` with shape `(1, 1, 32, 32)`, though I couldn't figure out exactly where it was. Is there a 1x1 conv2D in your code that maps 32 filters to 32 filters?

> From there, I can tell that the offending array has name kernel:0 with shape (1, 1, 32, 32), though I couldn't figure out exactly where it was. Is there a 1x1 conv2D in your code that maps 32 filters to 32 filters?

There are multiple 1x1 convs, yes. 

> Then run the tests with pytest --capture=no. This will break into a debugger at the point of failure, and you can step up to the calling frame with (u)p.

Could you elaborate a bit more here? I have added the `pdb` snippet into the model tester code. Then I ran `RUN_SLOW=1 python -m pytest --capture=no  tests/models/data2vec/test_modeling_tf_data2vec_vision.py`. I do get the pdb prompt and I get to `-> super().test_save_load()` as the oldest frame. 

@Rocketknight1 

@sayakpaul I stepped up to the frame of `dsid = dataset.make_new_dset(group, shape, dtype, data, name, **kwds)`. This let me inspect the variable `name` and the `group`, but I didn't understand `h5py` well enough to figure out the exact weight causing the issue.

@Rocketknight1 I looked into the layers with `kernel_size=1` and tried to fix their names to use something that's suffixed with identifiers. You can find the commit [here](https://github.com/sayakpaul/transformers/commit/8ccf88bf6bcc054307faf58e9ca2b21e04c6e60b).

It still didn't resolve the issue. The only potential suspect I could find is the following. There are two layers namely `classifier` in `TFData2VecVisionForSemanticSegmentation` that are added via `TFData2VecVisionUperHead` and `TFData2VecVisionFCNHead` respectively. 

Thoughts? 

Update:

With @Rocketknight1's help, I was able to resolve the current test failure (commit [here](https://github.com/sayakpaul/transformers/blob/fix/tf-data2vec-seg/src/transformers/models/data2vec/modeling_tf_data2vec_vision.py)). But I have run into two more failures which I am currently discussing with @Rocketknight1. He's on vacation. Once he gets back, hopefully, will be able to report back with updates.