TFBertModel much slower on GPU than BertModel
Author: jheinecke
Pull request: no

## Environment info

- `transformers` version: 4.8.2
- Platform: Linux-5.8.0-55-generic-x86_64-with-glibc2.31
- Python version: 3.9.5
- PyTorch version (GPU?): 1.9.0 (True)
- Tensorflow version (GPU?): 2.5.0 (True)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help

@LysandreJik, @Rocketknight1


## Information
Hello,
I have an GPU issue using BertModel with Pytorch 1.9.0 (OK) and Tensorflow 2.5.0 (very slow).
I'm trying to get vectors for each token of a sentence. The following minimal working example takes 15 seconds to process an (identical) sentence 2000 times (on a RTX 3090, cudatoolkit 11.1.74). GPU usage is about 30%

```
text = "this is a very simple sentence for which I need word piece vectors ."

def test_torch():
    import torch

    from transformers import BertModel, BertTokenizer

    tokenizer = BertTokenizer.from_pretrained("bert-base-multilingual-cased")
    model = BertModel.from_pretrained("bert-base-multilingual-cased")
    model = model.cuda()


    start = time.time()
    for x in range(2000):
        words = tokenizer.cls_token + text + tokenizer.sep_token # add [CLS] and [SEP] around sentence
        tokenids = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(words))
        #print(tokenids)
        token_tensors=torch.tensor([tokenids])
        with torch.no_grad():
            last_layer_features = model(token_tensors.cuda())[0]
            #print(last_layer_features)

    end = time.time()
    print("%d seconds" % (end-start))
```

whereas  using Tensorflow the same takes 130 seconds (and the GPU usage is about 4%):

```
def test_tf2():
    import tensorflow as tf
    from transformers import TFBertModel, BertTokenizer

    with tf.device("/gpu:0"):
        tokenizer = BertTokenizer.from_pretrained("bert-base-multilingual-cased")
        model = TFBertModel.from_pretrained("bert-base-multilingual-cased")
        #model = model.gpu() # does not exist

        start = time.time()
        for x in range(2000):
            words = tokenizer.cls_token + text + tokenizer.sep_token # add [CLS] and [SEP] around sentence
            tokenids = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(words))
            token_tensors = tf.convert_to_tensor([tokenids])
            #with torch.no_grad():
            last_layer_features = model(token_tensors.gpu(), training=False)[0]

        end = time.time()
        print("%d seconds" % (end-start))
```

So obviously I do not copy things correctly to the GPUs in the Tensorflow example. Does anyone have an idea what is wrong here?





Comments:
This is more of a TF issue than a Transformers one - eager TF can be quite inefficient, especially when you're repeatedly calling it with small inputs. In TF you can 'compile' a Python function or model call - this is analogous to using TorchScript in PyTorch, but is much more common and integrated into the main library - e.g. Keras almost always compiles models before running them. 

I tested your code with compilation and I got a speedup of 8x-10x, so I'd recommend using that. You can just add `model = tf.function(model)` before the loop, or for even higher performance on recent versions of TF you can use `model = tf.function(model, jit_compile=True)`. Alternatively, all of our models are Keras models, and if you run using the Keras API methods like fit() and predict(), all of the details of compilation and feeding data will be handled for you, and performance will usually be very good.

For more, I wrote a reddit post on the Keras approach with our library [here](https://www.reddit.com/r/MachineLearning/comments/ok81v4/n_tf_keras_and_transformers/), including links to example scripts.

You can also see a guide on using `tf.function` as a call or a decorator [here](https://www.tensorflow.org/guide/function).

Note that this code is a little bit sloppy - you'll ordinarily want to write the whole function that does what you want, and then call or decorate that function with `tf.function`, as explained in the TF docs. `tf.function(model)` works fine for this small example, but will cause problems if you then later want to use the model object for other things, such as using the Keras API with it.

Thanks for your reply ! I checked tf.function(model) and it works nicely. However my MWE was too minimal. In reality every sentence is different in wording and length. So I had to add a padding to make sure that all sentences have an equal (maximal) length, and now my "real" code is running as fast as pytorch. Apparently pytorch pads auto-magically, but in Tensorflow I have to do it manually. Anyway, it works! 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.