Issues running Transformer model example with estimator api
Hi @ali-raza-tariq ,

I am sorry to say that estimator API was not supported now and it will not work with TF2.x style code. Please refer to attached note below for same from the [source](https://www.tensorflow.org/guide/estimator).`

> Warning: Estimators are not recommended for new code. Estimators run [v1.Session](https://www.tensorflow.org/api_docs/python/tf/compat/v1/Session)-style code which is more difficult to write correctly, and can behave unexpectedly, especially when combined with TF 2 code. Estimators do fall under our [compatibility guarantees](https://tensorflow.org/guide/versions), but will receive no fixes other than security vulnerabilities. See the [migration guide](https://tensorflow.org/guide/migrate) for details.

However if you are passing non-python objects (like layers etc) you need to implement class method `from_config` to make them serializable. You may refer more details and a simple demo on how to implement `from_config` [here](https://www.tensorflow.org/guide/keras/serialization_and_saving#custom_objects). Since you are using layers in PatchEncoder class you must implement `from_config` class method.It may be of some help but still as the warning note suggested above it may or may not work with estimator API.

Thanks !



` @SuryanarayanaY  thanks for responding quickly! I understand there might be some limitation to what we can achieve with Estimator using TF2.x code, unfortunately I am in a situation where I need to evaluate various models and other models were using Estimator Api (so to ensure fair enough comparison - its ideal to use the same training api). I will try my best to get it working unless i know for sure that it cannot be done this way. Anyway - thanks to your response i was able to correctly implement the serialization and move past the above error. But even though i am not having issue saving/loading the model - it complains about missing variable values from the saved checkpoints.

this is how i am training the model:
```
model_est = keras.estimator.model_to_estimator(keras_model=model, model_dir='.', config=run_config)
train_spec = tf.estimator.TrainSpec(input_fn=train_input_fn, max_steps=100) 
eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn, throttle_secs=10)
tf.estimator.train_and_evaluate(model_est, train_spec, eval_spec)
``` 
and this is the error message i keep getting:
```
Traceback (most recent call last):
  File "train.py", line 266, in <module>
    history = run_experiment(vit_classifier)
  File "train.py", line 245, in run_experiment
    tf.estimator.train_and_evaluate(model_est, train_spec, eval_spec)
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow/python/util/deprecation.py", line 371, in new_func
    return func(*args, **kwargs)
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/training.py", line 503, in train_and_evaluate
    return executor.run()
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/training.py", line 644, in run
    return self.run_local()
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/training.py", line 741, in run_local
    self._estimator.train(
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 360, in train
    loss = self._train_model(input_fn, hooks, saving_listeners)
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1186, in _train_model
    return self._train_model_distributed(input_fn, hooks, saving_listeners)
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1247, in _train_model_distributed
    return self._actual_train_model_distributed(
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1360, in _actual_train_model_distributed
    return self._train_with_estimator_spec(estimator_spec, worker_hooks,
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1405, in _train_with_estimator_spec
    tf.compat.v1.train.warm_start(*self._warm_start_settings)
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow/python/training/warm_starting_util.py", line 467, in warm_start
    var_name_to_prev_var_name = _get_object_checkpoint_renames(
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow/python/training/warm_starting_util.py", line 397, in _get_object_checkpoint_renames
    raise ValueError(
ValueError: Attempting to warm-start from an object-based checkpoint, but found that the checkpoint did not contain values for all variables. The following variables were missing: {'kernel', 'bias', 'embeddings'}
```


 Hi @ali-raza-tariq ,

We are not supporting estimator API now. You may find similar issue from SO [here](https://stackoverflow.com/questions/63793180/tensorflow-checkpoint-variables-not-saved) for reference.

Please post the issue in Stack Overflow for support related issues where Community can also may be helpful.

Thank you! I understand ... thank you for the help! Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61393">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61393">No</a>

Hello everyone! I am trying to run an Image classification training example with Vision Transformer from keras examples (https://keras.io/examples/vision/image_classification_with_vision_transformer/). Everything ran perfectly when I ran it as it is but I started facing issues when i switched training from `model.fit()` to `tf.estimator.train_and_evaluate()` (ofcourse I made the appropriate changes to first convert model to estimator). From what I understand ... the problem lies with saving and reloading the model which is done by the estimator api. The model has custom classes:
```
class Patches(layers.Layer):
    def __init__(self, patch_size, **kwargs):
        super().__init__(**kwargs)
        self.patch_size = patch_size

    def call(self, images):
        batch_size = tf.shape(images)[0]
        patches = tf.image.extract_patches(
            images=images,
            sizes=[1, self.patch_size, self.patch_size, 1],
            strides=[1, self.patch_size, self.patch_size, 1],
            rates=[1, 1, 1, 1],
            padding="VALID",
        )
        patch_dims = patches.shape[-1]
        patches = tf.reshape(patches, [batch_size, -1, patch_dims])
        return patches

    ## personal addition
    def get_config(self):
        base_config = super().get_config()
        base_config.update({
            'patch_size': self.patch_size,
        })
        return base_config

class PatchEncoder(layers.Layer):
    def __init__(self, num_patches, projection_dim, **kwargs):
        super().__init__(**kwargs)
        self.num_patches = num_patches
        self.projection = layers.Dense(units=projection_dim)
        self.position_embedding = layers.Embedding(
            input_dim=num_patches, output_dim=projection_dim
        )

    def call(self, patch):
        positions = tf.range(start=0, limit=self.num_patches, delta=1)
        encoded = self.projection(patch) + self.position_embedding(positions)
        return encoded

    ## personal addition
    def get_config(self):
        base_config = super().get_config()
        base_config.update({
            'num_patches': self.num_patches,
            'projection': self.projection,
            'position_embedding': self.position_embedding
        })
        return base_config
```

From looking at some related issues, I found how we need to provide a `get_config()` method to save and reload the model with custom classes so I made small personal modifications but now its sort of giving me a different issue I am unable to understand. 

Error Log:
```
  warnings.warn(
x_train shape: (50000, 32, 32, 3) - y_train shape: (50000, 1)
x_test shape: (10000, 32, 32, 3) - y_test shape: (10000, 1)
WARNING:tensorflow:From train.py:225: RunConfig.__init__ (from tensorflow_estimator.python.estimator.run_config) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.keras instead.
/home/nearchus/.local/lib/python3.8/site-packages/keras/src/backend.py:452: UserWarning: `tf.keras.backend.set_learning_phase` is deprecated and will be removed after 2020-10-11. To update it, simply pass a True/False value to the `training` argument of the `__call__` method of your layer or model.
  warnings.warn(
Traceback (most recent call last):
  File "train.py", line 257, in <module>
    history = run_experiment(vit_classifier)
  File "train.py", line 231, in run_experiment
    model_est = keras.estimator.model_to_estimator(keras_model=model, model_dir='.', config=run_config)
  File "/home/nearchus/.local/lib/python3.8/site-packages/keras/src/estimator/__init__.py", line 376, in model_to_estimator_v2
    return keras_lib.model_to_estimator(
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/keras_lib.py", line 725, in model_to_estimator
    warm_start_path = _save_first_checkpoint(keras_model, custom_objects,
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/keras_lib.py", line 457, in _save_first_checkpoint
    model = _clone_and_build_model(ModeKeys.TRAIN, keras_model,
  File "/home/nearchus/.local/lib/python3.8/site-packages/tensorflow_estimator/python/estimator/keras_lib.py", line 230, in _clone_and_build_model
    clone = tf.compat.v2.keras.__internal__.models.clone_and_build_model(
  File "/home/nearchus/.local/lib/python3.8/site-packages/keras/src/models/cloning.py", line 806, in clone_and_build_model
    clone = clone_model(model, input_tensors=input_tensors)
  File "/home/nearchus/.local/lib/python3.8/site-packages/keras/src/models/cloning.py", line 539, in clone_model
    return _clone_functional_model(
  File "/home/nearchus/.local/lib/python3.8/site-packages/keras/src/models/cloning.py", line 222, in _clone_functional_model
    model_configs, created_layers = _clone_layers_and_model_config(
  File "/home/nearchus/.local/lib/python3.8/site-packages/keras/src/models/cloning.py", line 298, in _clone_layers_and_model_config
    config = functional.get_network_config(
  File "/home/nearchus/.local/lib/python3.8/site-packages/keras/src/engine/functional.py", line 1590, in get_network_config
    layer_config = serialize_layer_fn(layer)
  File "/home/nearchus/.local/lib/python3.8/site-packages/keras/src/models/cloning.py", line 295, in _copy_layer
    created_layers[layer.name] = layer_fn(layer)
  File "/home/nearchus/.local/lib/python3.8/site-packages/keras/src/models/cloning.py", line 52, in _clone_layer
    return layer.__class__.from_config(layer.get_config())
  File "train.py", line 108, in from_config
    return cls(**config)
TypeError: __init__() missing 1 required positional argument: 'projection_dim'
```
I thought it might be because of `PatchEncoder` class constructor has custom objects as argument - so i tried to do serialization/deserialization but to no vail. In any case, I would highly appreciate if someone can guide me as to where I am going wrong in this!  


  