[TensorFlow SAM] Internal operations in the prompt encoder fail during fine-tuning
Author: sayakpaul
Pull request: no

@merveenoyan and I are trying to create a fine-tuning notebook for the TensorFlow variant of [SAM](https://huggingface.co/docs/transformers/main/model_doc/sam). 

After compiling the model, when trying to run the actual fine-tuning, it leads to:

```
TypeError: in user code:

    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1284, in train_function  *
        return step_function(self, iterator)
    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1268, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/usr/local/lib/python3.10/dist-packages/keras/engine/training.py", line 1249, in run_step  **
        outputs = model.train_step(data)
    File "<ipython-input-10-3c1490d3fea1>", line 12, in train_step
        outputs = self.sam(
    File "/usr/local/lib/python3.10/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_filel5x0ne4l.py", line 37, in tf__run_call_with_unpacked_inputs
        retval_ = ag__.converted_call(ag__.ld(func), (ag__.ld(self),), dict(**ag__.ld(unpacked_inputs)), fscope)
    File "/tmp/__autograph_generated_filem4r4bspv.py", line 195, in tf__call
        (sparse_embeddings, dense_embeddings) = ag__.converted_call(ag__.ld(self).prompt_encoder, (), dict(batch_size=ag__.converted_call(ag__.ld(shape_list), (ag__.ld(image_embeddings),), None, fscope)[0], input_points=ag__.ld(input_points), input_labels=ag__.ld(input_labels), input_boxes=ag__.ld(input_boxes), input_masks=ag__.ld(input_masks)), fscope)
    File "/tmp/__autograph_generated_file4_u_db6c.py", line 90, in tf__call
        ag__.if_stmt(ag__.ld(input_boxes) is not None, if_body_3, else_body_3, get_state_3, set_state_3, ('batch_size', 'sparse_embeddings'), 2)
    File "/tmp/__autograph_generated_file4_u_db6c.py", line 68, in if_body_3
        box_embeddings = ag__.converted_call(ag__.ld(self)._embed_boxes, (ag__.ld(input_boxes),), None, fscope)
    File "/tmp/__autograph_generated_filehsxk3fhx.py", line 14, in tf___embed_boxes
        coords = ag__.converted_call(ag__.ld(tf).reshape, (ag__.ld(boxes), (ag__.ld(batch_size), ag__.ld(nb_boxes), 2, 2)), None, fscope)

    TypeError: Exception encountered when calling layer 'tf_sam_model' (type TFSamModel).
    
    in user code:
    
        File "/usr/local/lib/python3.10/dist-packages/transformers/modeling_tf_utils.py", line 1356, in run_call_with_unpacked_inputs  *
            return func(self, **unpacked_inputs)
        File "/usr/local/lib/python3.10/dist-packages/transformers/models/sam/modeling_tf_sam.py", line 1433, in call  *
            sparse_embeddings, dense_embeddings = self.prompt_encoder(
        File "/usr/local/lib/python3.10/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler  **
            raise e.with_traceback(filtered_tb) from None
        File "/tmp/__autograph_generated_file4_u_db6c.py", line 90, in tf__call
            ag__.if_stmt(ag__.ld(input_boxes) is not None, if_body_3, else_body_3, get_state_3, set_state_3, ('batch_size', 'sparse_embeddings'), 2)
        File "/tmp/__autograph_generated_file4_u_db6c.py", line 68, in if_body_3
            box_embeddings = ag__.converted_call(ag__.ld(self)._embed_boxes, (ag__.ld(input_boxes),), None, fscope)
        File "/tmp/__autograph_generated_filehsxk3fhx.py", line 14, in tf___embed_boxes
            coords = ag__.converted_call(ag__.ld(tf).reshape, (ag__.ld(boxes), (ag__.ld(batch_size), ag__.ld(nb_boxes), 2, 2)), None, fscope)
    
        TypeError: Exception encountered when calling layer 'prompt_encoder' (type TFSamPromptEncoder).
        
        in user code:
        
            File "/usr/local/lib/python3.10/dist-packages/transformers/models/sam/modeling_tf_sam.py", line 767, in call  *
                box_embeddings = self._embed_boxes(input_boxes)
            File "/usr/local/lib/python3.10/dist-packages/transformers/models/sam/modeling_tf_sam.py", line 726, in _embed_boxes  *
                coords = tf.reshape(boxes, (batch_size, nb_boxes, 2, 2))
        
            TypeError: Failed to convert elements of (None, None, 2, 2) to Tensor. Consider casting elements to a supported type. See https://www.tensorflow.org/api_docs/python/tf/dtypes for supported TF dtypes.
```

From the accompanying [Colab Notebook](https://colab.research.google.com/gist/sayakpaul/de59527f657d0461f46d9cb8c4a3884f/scratchpad.ipynb), one can check that there's nothing apparently off in the dataset we're passing to the trainer for fine-tuning:

```python
for sample in train_ds.take(2):
   for k in sample:
       print(k, sample[k].shape, isinstance(sample[k], tf.Tensor))
```

Leads to:

```bash
pixel_values (2, 3, 1024, 1024) True
original_sizes (2, 2) True
reshaped_input_sizes (2, 2) True
input_boxes (2, 1, 4) True
ground_truth_mask (2, 256, 256) True
pixel_values (2, 3, 1024, 1024) True
original_sizes (2, 2) True
reshaped_input_sizes (2, 2) True
input_boxes (2, 1, 4) True
ground_truth_mask (2, 256, 256) True
```

Anything we're missing out on? Cc: @Rocketknight1 

Comments:
Investigating now - when I was porting the model I got the feeling it didn't even really support fine-tuning! Is there a PyTorch notebook where fine-tuning works @sayakpaul?

Update: Some bits of the code were still using static shapes incorrectly! I've fixed it and I think your code sample should work now (there is a label shape issue, but I think that's not the model code's fault)

> Investigating now - when I was porting the model I got the feeling it didn't even really support fine-tuning! Is there a PyTorch notebook where fine-tuning works @sayakpaul?

Here you go: https://github.com/NielsRogge/Transformers-Tutorials/blob/master/SAM/Fine_tune_SAM_(segment_anything)_on_a_custom_dataset.ipynb

Thanks @Rocketknight1!

So, I had to transpose the predicted mask to have channels_last memory layout to make the loss computation work with Sparse Categorical Crossentropy:

```python
from tensorflow import keras

class SAMFineTuner(keras.Model):
    def __init__(self, sam, **kwargs):
        super().__init__(**kwargs)
        self.sam = sam

    def train_step(self, inputs):
        with tf.GradientTape() as tape:
            # Forward pass.
            outputs = self.sam(
                pixel_values=inputs["pixel_values"],
                input_boxes=inputs["input_boxes"],
                multimask_output=False
            )

            # Compute loss.
            predicted_masks = tf.squeeze(outputs.pred_masks, 1)
            predicted_masks = tf.transpose(predicted_masks, [0, 2, 3, 1])
            ground_truth_masks = tf.cast(inputs["ground_truth_mask"], tf.float32)
            loss = self.compiled_loss(tf.expand_dims(ground_truth_masks, 1), predicted_masks)

        # Optimize the model.
        trainable_vars = self.sam.trainable_variables
        grads = tape.gradient(loss, trainable_vars)
        self.optimizer.apply_gradients(zip(grads, trainable_vars))

        # Reporting.
        return {m.name: m.result() for m in self.metrics}
```

But we'll not be using SCCE anyway. 

Closing the issue, feel free to merge the PR :) 

Dear @sayakpaul, @Rocketknight1 and @merveenoyan,

Many thanks for your precious work.

I am following https://keras.io/examples/vision/sam/ to finetune HF [TFSamModel](https://huggingface.co/docs/transformers/main/model_doc/sam#transformers.TFSamModel). The training correctly works and my custom loss decreases in the tested epochs, but when I infer the test images, the predictions before or after the finetuning are the same. 
Therefore, I compare the `sam.get_weights()` before and after the finetune, and indeed they are equal. I expected differences, am I right?

I also tried to unident from `with tf.GradientTape() as tape:` 
```
# calculate loss over predicted and ground truth masks
loss = dice_loss(tf.expand_dims(ground_truth_masks, 1), predicted_masks)
# update trainable variables
trainable_vars = sam.trainable_variables
grads = tape.gradient(loss, trainable_vars)
optimizer.apply_gradients(zip(grads, trainable_vars))
return loss
```
as described in [Customizing what happens in fit()] (https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit#a_first_simple_example), but the predictions are the same and the weights too.

My second approach was to override [tf.keras.Model](https://www.tensorflow.org/api_docs/python/tf/keras/Model)  [train_step](https://github.com/keras-team/keras/blob/v2.14.0/keras/engine/training.py#L1128) as TFSamModel subclass it. 

The loss correctly decreases and the callbacks work but when the model weights before or after the finetune are the same.
```
Epoch 1/5
17/Unknown - 39s 441ms/step - loss: 0.9899 - mean_io_u_3: 0.4663
alpha updated to 0.9900000095367432
17/17 [==============================] - 40s 505ms/step - loss: 0.9899 - mean_io_u_3: 0.4663

Epoch 2/5
17/17 [==============================] - ETA: 0s - loss: 0.9799 - mean_io_u_3: 0.4663
alpha updated to 0.9800000190734863

17/17 [==============================] - 8s 442ms/step - loss: 0.9799 - mean_io_u_3: 0.4663
Epoch 3/5
17/17 [==============================] - ETA: 0s - loss: 0.9698 - mean_io_u_3: 0.4663
alpha updated to 0.9700000286102295

17/17 [==============================] - 7s 441ms/step - loss: 0.9698 - mean_io_u_3: 0.4663

Epoch 4/5
alpha updated to 0.9600000381469727
17/17 [==============================] - ETA: 0s - loss: 0.9599 - mean_io_u_3: 0.4663
17/17 [==============================] - 8s 441ms/step - loss: 0.9599 - mean_io_u_3: 0.4663

Epoch 5/5
alpha updated to 0.9500000476837158
17/17 [==============================] - ETA: 0s - loss: 0.9499 - mean_io_u_3: 0.4663
17/17 [==============================] - 8s 442ms/step - loss: 0.9499 - mean_io_u_3: 0.4663
*****
**
```

Do you have any ideas?

Thank you.

D

Leaving that one to @sayakpaul and @merveenoyan since they wrote that example, but please ping me if you think there are any issues in the underlying model!