Tencent / TPAT

TensorRT Plugin Autogen Tool
Apache License 2.0
365 stars 42 forks source link

KeyError int8 #37

Closed ngockhanh5110 closed 1 year ago

ngockhanh5110 commented 1 year ago

Model class

from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions

class CustomModel(GPT2LMHeadModel):
    def __init__(self, config):
        super(CustomModel, self).__init__(config)
        self.loss = torch.nn.CrossEntropyLoss()

    def forward(
        self,
        input_ids: Optional[torch.IntTensor] = None
    ) -> torch.FloatTensor:

        transformer_outputs = self.transformer(
            input_ids,
            past_key_values=None,
            attention_mask=None,
            token_type_ids=None,
            position_ids=None,
            head_mask=None,
            inputs_embeds=None,
            encoder_hidden_states=None,
            encoder_attention_mask=None,
            use_cache=None,
            output_attentions=None,
            output_hidden_states=None,
            return_dict=None,
        )
        hidden_states = transformer_outputs[0]

        lm_logits = self.lm_head(hidden_states)
        labels = input_ids
        shift_logits = lm_logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()

        loss = self.loss(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))        

        # return loss.reshape(-1,1)
        return CausalLMOutputWithCrossAttentions(
            loss=loss.reshape(-1,1),
            logits=None,
            past_key_values=None,
            hidden_states=None,
            attentions=None,
            cross_attentions=None,
        )

Command for conversion:

onnx2plugin(
    input_model_path ="./onnx_tpat/model.onnx", 
    output_model_path="./onnx_tpat/model.tpat.onnx", 
    # node_names="/loss/SoftmaxCrossEntropyLoss", 
    node_types = ["SoftmaxCrossEntropyLoss"], 
    plugin_name_dict={"SoftmaxCrossEntropyLoss": "tpat_softmax_cross_entropy"},
    dynamic_bs=False,
    # dynamic_bs=True, # if True, this operator support dynamic batchsize
    # min_bs=1,
    # opt_bs=64,
    # max_bs=100,
    )

I faced this error:

Couldn't find reusable plugin for node [/loss/SoftmaxCrossEntropyLoss](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f74656e736f7272742d6167632d6465746563746f722d72756e2d376531323764333238643637222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6563322d36332d33322d35322d3130302e65752d776573742d312e636f6d707574652e616d617a6f6e6177732e636f6d227d7d.vscode-resource.vscode-cdn.net/loss/SoftmaxCrossEntropyLoss)
Start auto-tuning!
Compile...
[/tmp/tuning.log](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f74656e736f7272742d6167632d6465746563746f722d72756e2d376531323764333238643637222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6563322d36332d33322d35322d3130302e65752d776573742d312e636f6d707574652e616d617a6f6e6177732e636f6d227d7d.vscode-resource.vscode-cdn.net/tmp/tuning.log) does not exist!

Running...
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[2], line 1
----> 1 onnx2plugin(
      2     input_model_path ="[./onnx_tpat/model.onnx](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f74656e736f7272742d6167632d6465746563746f722d72756e2d376531323764333238643637222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6563322d36332d33322d35322d3130302e65752d776573742d312e636f6d707574652e616d617a6f6e6177732e636f6d227d7d.vscode-resource.vscode-cdn.net/root/working/onnx_tpat/model.onnx)", 
      3     output_model_path="[./onnx_tpat/model.tpat.onnx](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f74656e736f7272742d6167632d6465746563746f722d72756e2d376531323764333238643637222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6563322d36332d33322d35322d3130302e65752d776573742d312e636f6d707574652e616d617a6f6e6177732e636f6d227d7d.vscode-resource.vscode-cdn.net/root/working/onnx_tpat/model.tpat.onnx)", 
      4     # node_names="[/loss/SoftmaxCrossEntropyLoss](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f74656e736f7272742d6167632d6465746563746f722d72756e2d376531323764333238643637222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6563322d36332d33322d35322d3130302e65752d776573742d312e636f6d707574652e616d617a6f6e6177732e636f6d227d7d.vscode-resource.vscode-cdn.net/loss/SoftmaxCrossEntropyLoss)", 
      5     node_types = ["SoftmaxCrossEntropyLoss"], 
      6     plugin_name_dict={"SoftmaxCrossEntropyLoss": "tpat_softmax_cross_entropy"},
      7     dynamic_bs=False,
      8     # dynamic_bs=True, # if True, this operator support dynamic batchsize
      9     # min_bs=1,
     10     # opt_bs=64,
     11     # max_bs=100,
     12     )

File [/workspace/TPAT/python/onnx_to_plugin.py:196](https://vscode-remote+attached-002dcontainer-002b7b22636f6e7461696e65724e616d65223a222f74656e736f7272742d6167632d6465746563746f722d72756e2d376531323764333238643637222c2273657474696e6773223a7b22686f7374223a227373683a2f2f6563322d36332d33322d35322d3130302e65752d776573742d312e636f6d707574652e616d617a6f6e6177732e636f6d227d7d.vscode-resource.vscode-cdn.net/workspace/TPAT/python/onnx_to_plugin.py:196), in onnx2plugin(input_model_path, output_model_path, node_names, node_types, plugin_name_dict, dynamic_bs, min_bs, max_bs, opt_bs)
    194         os.remove(dy_input_model)
    195 else:
--> 196     onnx_name_mapping_trt_plugin = generate_plugin_library(
    197         input_model_path, nodes, plugin_name_dict 
    198     )
    199 print("Onnx_name_mapping_trt_plugin: {}".format(onnx_name_mapping_trt_plugin))
    200 OnnxModified(
    201     input_model_path, output_model_path, nodes, onnx_name_mapping_trt_plugin
...
    352             )
    353     input_slot_dict[idx] = self._input_dict[str(i)]
    354 if len(self._allocate_global_memory) != 0:

KeyError: 'int8'
ngockhanh5110 commented 1 year ago

Have another approach. Thanks