From 6912265711fdf29e537e35194552bfa2c31b72bc Mon Sep 17 00:00:00 2001 From: Abel Date: Tue, 7 Jul 2020 11:32:29 +0200 Subject: [PATCH] Make T5 compatible with ONNX (#5518) * Default decoder inputs to encoder ones for T5 if neither are specified. * Fixing typo, now all tests are passing. * Changing einsum to operations supported by onnx * Adding a test to ensure T5 can be exported to onnx op>9 * Modified test for onnx export to make it faster * Styling changes. * Styling changes. * Changing notation for matrix multiplication Co-authored-by: Abel Riboulot --- src/transformers/modeling_t5.py | 21 ++++++++++++++++----- tests/test_modeling_t5.py | 10 ++++++++++ 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/src/transformers/modeling_t5.py b/src/transformers/modeling_t5.py index ce044efe6c..76cf176c5e 100644 --- a/src/transformers/modeling_t5.py +++ b/src/transformers/modeling_t5.py @@ -358,7 +358,10 @@ class T5Attention(nn.Module): else: present_key_value_state = (None,) - scores = torch.einsum("bnqd,bnkd->bnqk", q, k) # (bs, n_heads, qlen, klen) + # (bs, n_heads, qlen, klen) + scores = torch.matmul( + q, k.transpose(3, 2) + ) # equivalent of torch.einsum("bnqd,bnkd->bnqk", q, k), compatible with onnx op>9 if position_bias is None: if not self.has_relative_attention_bias: @@ -818,7 +821,8 @@ T5_INPUTS_DOCSTRING = r""" Provide for sequence to sequence training. T5 uses the pad_token_id as the starting token for decoder_input_ids generation. If `decoder_past_key_value_states` is used, optionally only the last `decoder_input_ids` have to be input (see `decoder_past_key_value_states`). To know more on how to prepare :obj:`decoder_input_ids` for pre-training take a look at - `T5 Training <./t5.html#training>`__. + `T5 Training <./t5.html#training>`__. If decoder_input_ids and decoder_inputs_embeds are both None, + decoder_input_ids takes the value of input_ids. decoder_attention_mask (:obj:`torch.BoolTensor` of shape :obj:`(batch_size, tgt_seq_len)`, `optional`, defaults to :obj:`None`): Default behavior: generate a tensor that ignores pad tokens in decoder_input_ids. Causal mask will also be used by default. decoder_past_key_value_states (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): @@ -837,7 +841,8 @@ T5_INPUTS_DOCSTRING = r""" Optionally, instead of passing :obj:`decoder_input_ids` you can choose to directly pass an embedded representation. If `decoder_past_key_value_states` is used, optionally only the last `decoder_inputs_embeds` have to be input (see `decoder_past_key_value_states`). This is useful if you want more control over how to convert `decoder_input_ids` indices into associated vectors - than the model's internal embedding lookup matrix. + than the model's internal embedding lookup matrix. If decoder_input_ids and decoder_inputs_embeds are both None, + decoder_inputs_embeds takes the value of inputs_embeds. head_mask: (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`, defaults to :obj:`None`): Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``: @@ -934,7 +939,7 @@ class T5Model(T5PreTrainedModel): >>> model = T5Model.from_pretrained('t5-small') >>> input_ids = tokenizer.encode("Hello, my dog is cute", return_tensors="pt") # Batch size 1 - >>> outputs = model(input_ids=input_ids, decoder_input_ids=input_ids) + >>> outputs = model(input_ids=input_ids) >>> last_hidden_states = outputs[0] # The last hidden-state is the first element of the output tuple """ @@ -953,6 +958,12 @@ class T5Model(T5PreTrainedModel): hidden_states = encoder_outputs[0] + # If the model is only provided with either input_ids or inputs_embeds, + # use them as the inputs of the decoder. self.encoder checks for input_ids XOR inputs_embeds + if (decoder_input_ids is None) and (decoder_inputs_embeds is None): + decoder_input_ids = input_ids + decoder_inputs_embeds = inputs_embeds + # If decoding with past key value states, only the last tokens # should be given as an input if decoder_past_key_value_states is not None: @@ -1076,7 +1087,7 @@ class T5ForConditionalGeneration(T5PreTrainedModel): >>> tokenizer = T5Tokenizer.from_pretrained('t5-small') >>> model = T5ForConditionalGeneration.from_pretrained('t5-small') >>> input_ids = tokenizer.encode("Hello, my dog is cute", return_tensors="pt") # Batch size 1 - >>> outputs = model(input_ids=input_ids, decoder_input_ids=input_ids, labels=input_ids) + >>> outputs = model(input_ids=input_ids, labels=input_ids) >>> loss, prediction_scores = outputs[:2] >>> tokenizer = T5Tokenizer.from_pretrained('t5-small') diff --git a/tests/test_modeling_t5.py b/tests/test_modeling_t5.py index a75e46a8fd..6330ef18de 100644 --- a/tests/test_modeling_t5.py +++ b/tests/test_modeling_t5.py @@ -351,6 +351,16 @@ class T5ModelTest(ModelTesterMixin, unittest.TestCase): model = T5Model.from_pretrained(model_name) self.assertIsNotNone(model) + def test_export_to_onnx(self): + import tempfile + + config_and_inputs = self.model_tester.prepare_config_and_inputs() + model = T5Model(config_and_inputs[0]) + with tempfile.TemporaryDirectory() as tmpdirname: + torch.onnx.export( + model, config_and_inputs[1], f"{tmpdirname}/t5_test.onnx", export_params=True, opset_version=9, + ) + @require_torch class T5ModelIntegrationTests(unittest.TestCase):