Model's internal loss is different than calculated loss at tf.keras.metric.Metric
Author: uunal
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.22.1
- Platform: Linux-4.19.0-21-cloud-amd64-x86_64-with-glibc2.28
- Python version: 3.9.5

### Who can help?

@Rocketknight1

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Using tf.keras.Model version of gpt2, and training as casual language modelling. 
```
gpt2_model.compile(..., metrics=[PerplexityMetric()]) #using inner loss (hf_compute_loss)
gpt2_model.fit(...)
```
The problem occurs when metric is calculated. I did not come up with an idea to use inner loss of transformers in tf.keras.metric.Metric (if anyone knows how it's good to rid of duplicate calculation) . So calculating loss as written in def hf_compute_loss in TFCausalLanguageModelingLoss. PerplexityMetric as:

```
class PerplexityMetric(tf.keras.metrics.Metric):    
    def __init__(self, name='perplexity', **kwargs):
      super(PerplexityMetric, self).__init__(name=name, **kwargs)
      self.cross_entropy = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
      self.perplexity = self.add_weight(name='tp', initializer='zeros')
	
    def _calculate_perplexity(self, real, pred):	       	 			
        unmasked_loss = self.cross_entropy(tf.nn.relu(real), pred)
        # make sure only labels that are not equal to -100 affect the loss
        loss_mask = tf.cast(real != -100, dtype=unmasked_loss.dtype)
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(
            loss_mask
        )
        loss_ = tf.reshape(reduced_masked_loss, (1,))        
        perplexity = tf.math.exp(loss_[-1])        
        return perplexity 

    def update_state(self, y_true, y_pred, sample_weight=None):           
        perplexity = self._calculate_perplexity(y_true, y_pred)	        	
        self.perplexity.assign(perplexity)
        
    def result(self):
      return self.perplexity

    def reset_state(self):
      # reset at the start of each epoch.
      self.perplexity.assign(0.0)
```


### Expected behavior

Both inner loss calculation and the loss calculation in the metric are same. What I expect is that calculated losses are the same in training output. Last float within brackets is loss calculated in PerplexityMetric() (_loss__).

```
1/7 [===>..........................] - ETA: 2s - loss: 8.0234 - perplexity: 1708.3318[7.44036245]
2/7 [=======>......................] - ETA: 2s - loss: 8.1277 - perplexity: 1703.3676[7.78889942]
3/7 [===========>..................] - ETA: 1s - loss: 8.2354 - perplexity: 2413.6597[7.75492764]
4/7 [================>.............] - ETA: 1s - loss: 8.3068 - perplexity: 2333.0405[7.91714859]
5/7 [====================>.........] - ETA: 0s - loss: 8.3837 - perplexity: 2743.9358[7.48942709]
6/7 [========================>.....] - ETA: 0s - loss: 8.3680 - perplexity: 1789.0269[7.8069911]
```
I did not understand the difference between loss values. Is there any wrong calculation in PerplexityMetric or is this a buggy behaviour?

Comments:
Hi @uunal, I don't think this is a bug in `transformers`. The main issue is that Keras displays the loss as the average for the whole epoch so far, but your Metric only displays the loss for the most recent batch. This is because you overwrite it with each batch. A better approach might be to store `self.total_loss` and `self.num_batches` and then in `result()` you could return `tf.math.exp(self.total_loss / self.num_batches)`.

I'm going to close this issue for now, but if you believe you've identified a bug in `transformers`, feel free to re-open it!

Thank you for the quick response and your suggestion. I understood the difference in losses.
If anyone interested later on, adding the updated metric code below:

```
class PerplexityMetric(tf.keras.metrics.Metric):
    def __init__(self, name="perplexity", **kwargs):
        super(PerplexityMetric, self).__init__(name=name, **kwargs)
        self.cross_entropy = tf.keras.losses.SparseCategoricalCrossentropy(
            from_logits=True, reduction=tf.keras.losses.Reduction.NONE
        )
        self.total_loss = self.add_weight(name="pl", shape=(), initializer="zeros")
        self.num_batches = self.add_weight(
            name="nb", shape=(), initializer="zeros"
        )

    def _calculate_loss(self, real, pred):
        unmasked_loss = self.cross_entropy(tf.nn.relu(real), pred)
        # make sure only labels that are not equal to -100 affect the loss
        loss_mask = tf.cast(real != -100, dtype=unmasked_loss.dtype)
        masked_loss = unmasked_loss * loss_mask
        reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask)
        loss_ = tf.reshape(reduced_masked_loss, (1,))
        return loss_[-1]

    def update_state(self, y_true, y_pred, sample_weight=None):
        loss_ = self._calculate_loss(y_true, y_pred)
        # update state variables
        self.num_batches.assign_add(1)
        self.total_loss.assign_add(loss_)

    def result(self):
        # perplexity: subword-level       
        return tf.math.exp(self.total_loss / self.num_batches)

    def reset_state(self):
        # reset at the start of each epoch.
        self.total_loss.assign(0.0)
        self.num_batches.assign(0)
```