Exporting Fine tuned T5ForConditionalGeneration model to TF-Serving using ONNX
Author: sekharvth
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.9.1
- Platform: Linux-5.4.0-1049-gcp-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.9.0 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help
@patrickvonplaten, @patil-suraj 

## Information

Model I am using (Bert, XLNet ...): T5

The problem arises when using:
* my own modified scripts: 
I use a fine tuned version of t5 (fine tuned using Huggingfqce and PyTorch), trained on a custom dataset for summarization. Since PyTorch Serving is no longer an option because of unrelated reasons, I require TF-Serving for a production optimized setting. I'm using the ONNX pipeline detailed here : https://huggingface.co/transformers/serialization.html#converting-an-onnx-model-using-the-transformers-onnx-package, with necessary changes to the paths. 

When I serve this model and do inference, it seems the model being loaded isn't the fine tuned one, as it gives output of the following nature : `In In In In In auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf auf`

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: Summarization

## To reproduce
Steps to reproduce the behavior:
```
def load_ckp(checkpoint_fpath, model, optimizer):
    checkpoint = torch.load(checkpoint_fpath, map_location=torch.device('cpu'))
    model.load_state_dict(checkpoint['state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    return model, optimizer, checkpoint['epoch']

import os

tokenizer = T5Tokenizer.from_pretrained("t5-base")
model = T5ForConditionalGeneration.from_pretrained("t5-base")
from torch import cuda
device = 'cuda' if cuda.is_available() else 'cpu'
model = model.to(device)
optimizer = torch.optim.Adam(params =  model.parameters(), lr=1e-4)

ckp_path = '/checkpoint_dir/checkpoint.pt'
model, optimizer, start_epoch = load_ckp(ckp_path, model, optimizer)

model.save_pretrained('fine-tuned') ##saves it to pytorch_model.bin format and config.json, which is needed for onnx
tokenizer.save_pretrained('fine-tuned')

tokenizer = T5Tokenizer.from_pretrained("fine-tuned")
model = T5ForConditionalGeneration.from_pretrained("fine-tuned")

==== COMMAND LINE====

python -m transformers.onnx --model=fine-tuned onnx/t5-tf-serving/
```
At this point, I get the following warning : `Some weights of the model checkpoint at fine-tuned were not used when initializing T5Model: ['lm_head.weight']`, 
but the process completes with the following message : `All good, model saved at: onnx/t5-tf-serving/model.onnx`.

Post this I use `onnx-tf convert -i onnx/t5-tf-serving/model.onnx -o output.pb` to get the corresponding Tensorflow SavedModel, and use standard docker based procedure for deploying it with TF-Serving.


## Expected behavior

I'm able to serve the model using the proper request formats, but the outputs are way off, as shown above. I'm guessing it has to do with the warning message that was displayed when converting the pytorch model to onnx. 
Fwiw, I tested out normal inference on the .bin formatted pytorch model that was obtained using the `model.save_pretrained('fine-tuned')` function, and it was generating expected outputs. 

Can you please suggest workarounds? @patrickvonplaten, @patil-suraj 


Comments:
Also pinging @mfuntowicz here as the expert for everything onnx related. It might be the case that the `lm_head.weight` are not correctly ported to ONNX here. 

Could you try to do the following after training

After line: 
```
model, optimizer, start_epoch = load_ckp(ckp_path, model, optimizer)
```
add those two lines 

```
model.config.tie_word_embeddings =False  # sets `config.tie_word_embeddings=False`
model.lm_head.weight = torch.nn.Parameter(model.shared.weight.T.clone())  # clones the correct parameters
```
and try if this works. Not at all sure, but this should at least remove the warning

Hey @patrickvonplaten, thanks for getting back.

I tried running the pipeline with the additional lines of code you'd provided. It gave me a shape mismatch initially in the line where the `model.lm_head.weight` parameter is set, but I was able to get around that by removing the `.T` in the RHS. 

Even after doing this, the `lm_head.weight not initialized` warning still appeared, but I decided to try out the remainder anyway. Curiously, I did notice that the `config.json` file had an additional key-value pair this time, specifically `'decoder_start_token_id': 0`, which I'm hoping was a step in the right direction.

I converted the onnx file that was generated to the TF SavedModel format with the same command in my original post, and inspected the inputs and outputs to the model with the `saved_model_cli show` command. The outputs of the same are as follows : 

```
signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['attention_mask'] tensor_info:
        dtype: DT_INT64
        shape: (-1, -1)
        name: serving_default_attention_mask:0
    inputs['decoder_attention_mask'] tensor_info:
        dtype: DT_INT64
        shape: (-1, 1)
        name: serving_default_decoder_attention_mask:0
    inputs['decoder_input_ids'] tensor_info:
        dtype: DT_INT64
        shape: (-1, 1)
        name: serving_default_decoder_input_ids:0
    inputs['input_ids'] tensor_info:
        dtype: DT_INT64
        shape: (-1, -1)
        name: serving_default_input_ids:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['output_0'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, -1, 768)
        name: StatefulPartitionedCall:0
    outputs['output_1'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, -1, 768)
        name: StatefulPartitionedCall:1
```

The outputs section indicates that these are the outputs of a couple of hidden layers which has 768 units (I'm guessing it's the final encoder/decoder hidden states), but what is actually required is the softmax distribution over ~32k tokens for summarization. 

Just out of curiosity, I did try exporting the fine tuned PyTorch model through the standard `torch.onnx.export` pipeline, with the following command:
`torch.onnx.export(model, (model.dummy_inputs['input_ids'], model.dummy_inputs['decoder_input_ids'], model.dummy_inputs['decoder_attention_mask']), 't5-for-serve.onnx')`, and it exited successfully with just a `TracerWarning`.

On converting this to the TF SavedModel format, I get these input and output signatures : 

```
signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['0'] tensor_info:
        dtype: DT_INT64
        shape: (3, 5)
        name: serving_default_0:0
    inputs['2'] tensor_info:
        dtype: DT_INT64
        shape: (3, 5)
        name: serving_default_2:0
    inputs['attention_mask.1'] tensor_info:
        dtype: DT_INT64
        shape: (3, 5)
        name: serving_default_attention_mask.1:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['output_0'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 5, 32128)
        name: PartitionedCall:0
    outputs['output_1'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:1
    outputs['output_10'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:2
    outputs['output_11'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:3
    outputs['output_12'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:4
    outputs['output_13'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:5
    outputs['output_14'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:6
    outputs['output_15'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:7
    outputs['output_16'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:8
    outputs['output_17'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:9
    outputs['output_18'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:10
    outputs['output_19'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:11
    outputs['output_2'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:12
    outputs['output_20'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:13
    outputs['output_21'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:14
    outputs['output_22'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:15
    outputs['output_23'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:16
    outputs['output_24'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:17
    outputs['output_25'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:18
    outputs['output_26'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:19
    outputs['output_27'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:20
    outputs['output_28'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:21
    outputs['output_29'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:22
    outputs['output_3'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:23
    outputs['output_30'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:24
    outputs['output_31'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:25
    outputs['output_32'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:26
    outputs['output_33'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:27
    outputs['output_34'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:28
    outputs['output_35'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:29
    outputs['output_36'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:30
    outputs['output_37'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:31
    outputs['output_38'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:32
    outputs['output_39'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:33
    outputs['output_4'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:34
    outputs['output_40'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:35
    outputs['output_41'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:36
    outputs['output_42'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:37
    outputs['output_43'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:38
    outputs['output_44'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:39
    outputs['output_45'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:40
    outputs['output_46'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:41
    outputs['output_47'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:42
    outputs['output_48'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:43
    outputs['output_49'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 5, 768)
        name: PartitionedCall:44
    outputs['output_5'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:45
    outputs['output_6'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:46
    outputs['output_7'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:47
    outputs['output_8'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:48
    outputs['output_9'] tensor_info:
        dtype: DT_FLOAT
        shape: (3, 12, 5, 64)
        name: PartitionedCall:49
```

As can be seen, the inputs and outputs differ significantly. The number of inputs have gone down by 1, compared to the `transformers.onnx` export having suggestions from @patrickvonplaten. Plus, the sizes of the inputs and outputs are fixed, which I'm guessing can be overriden with the `dynamic_axes` parameter in the export step. 

But the first output in this does have what appears to be the softmax distribution over the ~32k tokens. 

Is this the right approach @patrickvonplaten @mfuntowicz ?


Hmm ok :-/ Think we'll have to wait here until @mfuntowicz is back from vacation I'm afraid...he knows better how to debug onnx + tf

Sure @patrickvonplaten thanks for the quick turnaround though.

Hey @patrickvonplaten, just checking in to see if @mfuntowicz is back :)

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.

Hey @patrickvonplaten @mfuntowicz , any updates on this?

Hi @sekharvth,
It seems that you are interested in doing seq2seq language modeling (`T5ForConditionalGeneration`), but you do not specify that when exporting to ONNX, so what seems to happen is that you are only exporting a `T5Model`, not a `T5ForConditionalGeneration`, hence the issue.

For each architecture, you can find which feature is supported for the ONNX export [here](https://github.com/huggingface/transformers/blob/master/src/transformers/onnx/features.py#L60). 

The command line utilty takes a feature parameter that defaults to `default` when left unspecified, in your case you want `seq2seq-lm`, could you try this and tell me if this solves the issue:

```
python -m transformers.onnx --feature="seq2seq-lm" --model=fine-tuned onnx/t5-tf-serving/ output_dir
```





Hey @michaelbenayoun .

Aaah, there's a param to be set then. I'll try this out and let you know :)

Thank you!

Hey @michaelbenayoun , thank you so much for pointing out the solution, it works like a charm :)

Side note : If anyone is else is trying to serve a generative model, you need to make calls to the served model repeatedly with the latest decoded array, until the stop token is generated (much like the usual inference logic). Takeaway being that from limited initial testing, serving the model doesn't necessarily speed up the inference process. 

Closing this issue :)