Segformer, can't save checkpoint in saved_model format
Author: joihn
Pull request: no
Labels: name: bug

Thanks for this repo ! 
###  Issue description:
1) initalise your segformer model
2) try to save it.

### Error:
```
    File "/home/maxime/anaconda3/envs/tf/lib/python3.9/site-packages/transformers/models/segformer/modeling_tf_segformer.py", line 547, in serving  *
        output = self.call(inputs)
    File "/home/maxime/anaconda3/envs/tf/lib/python3.9/site-packages/transformers/models/segformer/modeling_tf_segformer.py", line 753, in call  *
        batch_size = shape_list(encoder_hidden_states[-1])[0]

    KeyError: -1
```

### Minimum reproducting code:

```
import os
import tensorflow as tf
from transformers import TFSegformerForSemanticSegmentation


model = TFSegformerForSemanticSegmentation.from_pretrained(
    "nvidia/mit-b0",
    num_labels=2,
    id2label={1:"1", 2:"2"},
    ignore_mismatched_sizes=True,  # Will ensure the segmentation specific components are reinitialized.

)

model.summary(line_length=250)
tf.keras.models.save_model(
    model, os.path.join("/tmp", model.name), include_optimizer=False
)
```

### Expected behavior
The checkpoint should save correctly

### System Info
Ubuntu 20.04, Python 3.9, TF 2.9.1, Nvidia Titan

### Who can help?

@sayakpaul  @NielsRogge 

Comments:
Have you tried using `model.save_pretrained()`?

Cc: @amyeroberts 

`model.save_pretrained()` indeed works, thanks :)

Shall we close this issue, or is this model also supposed to be compatible with keras saving method? (their callback is widely used, https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/ModelCheckpoint)

@joihn Glad to hear you were able to save with `save_pretrained` and thanks for responding so quickly @sayakpaul!

I'll defer the issue to our TF gurus @Rocketknight1 @gante regarding compatibility with keras saving. 

I see two possible options for the time being:

* Implement a custom callback to use `save_pretrained()`. Shouldn't differ too much from the `ModelCheckpoint` callback. 
*  You can refer to [this notebook](https://github.com/huggingface/notebooks/blob/main/examples/image_classification-tf.ipynb) that makes use of `PushToHubCallback()` and achieves a similar result as `ModelCheckpoint` barring some differences. 

I did some digging, this issue seems to appears when keras function [save_model] https://www.tensorflow.org/api_docs/python/tf/keras/models/save_model) parmeters is `save_traces=True`.

It's worth noting that hugging face `saved_pretrained(filepath, saved_model=True)` also crash with the same error.

also related to https://github.com/huggingface/transformers/issues/13742


Hi @joihn - there are some general difficulties when saving Hugging Face models as SavedModel. This is a general issue with any model where the model and layers are implemented by subclassing in Keras - SavedModel doesn't really have a good way to completely save and load those models (although you can save one or more model traces through SavedModel, this isn't usually what people want unless they're trying to export to TFLite or something!)

Instead, we recommend that users save weights only, and if they want to save the entire model, to use the `save_pretrained` method, which will save the weights along with a config that will make it loadable with the `load_pretrained` method.

Concretely, this means doing the following things:

1) When using `ModelCheckpoint`, set `save_weights_only` to `True`.
2) Replace `model.save` with either `model.save_weights` or `model.save_pretrained`

Perfect, thanks for the info :) 

`model.saved_pretrained` in fact provides the flag `saved_model=True`  in order to save the model as SavedModel besides `h5` weights and `config.json`. However I find that `include_optimizer` is set to `False`, which means the SavedModel cannot be resumed to continue the training. Is there a strong reason for this hard coding ?

```        if saved_model:
            # If `torch_dtype` is in the config with a torch dtype class as the value, we need to change it to string.
            # (Although TF doesn't care about this attribute, we can't just remove it or set it to `None`.)
            if getattr(self.config, "torch_dtype", None) is not None and not isinstance(self.config.torch_dtype, str):
                self.config.torch_dtype = str(self.config.torch_dtype).split(".")[1]
            if signatures is None:
                serving_default = self.serving.get_concrete_function(self.input_signature)
                if any(spec.dtype == tf.int32 for spec in self.input_signature.values()):
                    int64_spec = {
                        key: tf.TensorSpec(
                            shape=spec.shape, dtype=tf.int64 if spec.dtype == tf.int32 else spec.dtype, name=spec.name
                        )
                        for key, spec in self.input_signature.items()
                    }
                    int64_serving = self.serving.get_concrete_function(int64_spec)
                    signatures = {"serving_default": serving_default, "int64_serving": int64_serving}
                else:
                    signatures = serving_default
            saved_model_dir = os.path.join(save_directory, "saved_model", str(version))
            self.save(saved_model_dir, include_optimizer=False, signatures=signatures)
            logger.info(f"Saved model created in {saved_model_dir}")
```