Segformer, can't save checkpoint in saved_model format

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

[Rocketknight1]

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

[amyeroberts]

@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.