Commit 7bbe528d authored by Haoran Li's avatar Haoran Li Committed by Facebook Github Bot
Browse files

fixes on bi-transformer onnx 

Summary: replace dynamic index put with copying and creating a new tensor

Reviewed By: wanchaol

Differential Revision: D13244573

fbshipit-source-id: 909f7913ad579ed035f29bb52321ff01e09a2c60
parent 866d0d2e

fairseq/modules/character_token_embedder.py

+16 −2
Original line number Diff line number Diff line
@@ -32,6 +32,7 @@ class CharacterTokenEmbedder(torch.nn.Module): 
    ):

        super(CharacterTokenEmbedder, self).__init__()



        self.onnx_trace = False

        self.embedding_dim = word_embed_dim

        self.max_char_len = max_char_len

        self.char_embeddings = nn.Embedding(257, char_embed_dim, padding_idx=0)
@@ -58,6 +59,9 @@ class CharacterTokenEmbedder(torch.nn.Module): 


        self.reset_parameters()



    def prepare_for_onnx_export_(self):

        self.onnx_trace = True



    def set_vocab(self, vocab, max_char_len):

        word_to_char = torch.LongTensor(len(vocab), max_char_len)
@@ -101,7 +105,11 @@ class CharacterTokenEmbedder(torch.nn.Module): 
            pads = chars[:, 0].eq(CHAR_PAD_IDX)

            eos = chars[:, 0].eq(CHAR_EOS_IDX)

            if eos.any():

                if self.onnx_trace:

                    chars = torch.where(eos.unsqueeze(1), chars.new_zeros(1), chars)

                else:

                    chars[eos] = 0



            unk = None

        else:

            flat_words = input.view(-1)
@@ -111,12 +119,18 @@ class CharacterTokenEmbedder(torch.nn.Module): 
            unk = flat_words.eq(self.vocab.unk())



        word_embs = self._convolve(chars)

        if self.onnx_trace:

            if pads.any():

                word_embs = torch.where(pads.unsqueeze(1), word_embs.new_zeros(1), word_embs)

            if eos.any():

                word_embs = torch.where(eos.unsqueeze(1), self.symbol_embeddings[self.eos_idx], word_embs)

            if unk is not None and unk.any():

                word_embs = torch.where(unk.unsqueeze(1), self.symbol_embeddings[self.unk_idx], word_embs)

        else:

            if pads.any():

                word_embs[pads] = 0



            if eos.any():

                word_embs[eos] = self.symbol_embeddings[self.eos_idx]



            if unk is not None and unk.any():

                word_embs[unk] = self.symbol_embeddings[self.unk_idx]