Help with Using TFXLNet on custom embeddings
Author: andrewlee98
Pull request: no
Labels: name: wontfix

# 🐛 Bug

## Information

Hi, I am working on implementing the multimodal bitransformer: https://arxiv.org/pdf/1909.02950.pdf
I have already gotten this working using your implementation of TFBertModel, but I want to try using TFXLNetModel in place of BERT to see if it makes an improvement. What I've done for using TFBertModel is extract the word_embeddings and pass the word embeddings with token_type_ids = 0 and image embeddings with token_type_ids = 1. 

## To reproduce

I have the following definition of my model

`
class BERT(transformers.TFXLNetModel):
    def __init__(self, config, *inputs, **kwargs):
        super(BERT, self).__init__(config, *inputs, **kwargs)
        self.call = tf.function(self.call)


class MyModel(tf.keras.Model):
    
    def __init__(self):
        super(MyModel, self).__init__()
        self.resnet = tf.keras.applications.ResNet152V2(include_top=False, 
                                                             weights='imagenet', 
                                                             input_shape=(224, 224, 3))
        self.bert = BERT.from_pretrained('xlnet-base-cased')
        
        self.text_embedding = self.bert.get_input_embeddings().weights[0]
         
        self.pooling = layers.AveragePooling2D(pool_size=(2, 2), padding='same')
        self.reshape = layers.Reshape((4 * 4, 2048)) # 4 is from 7//2 + 1
        self.W_ns = [layers.Dense(self.bert.config.hidden_size) for _ in range(self.reshape.target_shape[0])]
        
        self.concat = layers.Concatenate(axis=1)
        
        self.dropout = layers.Dropout(0.1)
        self.denseout = layers.Dense(1, activation='sigmoid')

    def call(self, inputs):
        text, image = inputs
        
        # handle image
        image = tf.keras.applications.resnet_v2.preprocess_input(image)
        image_emb = self.resnet(image)
        image_emb = self.pooling(image_emb)
        image_emb = self.reshape(image_emb)
        image_embeds = [self.W_ns[i](image_emb[:, i]) for i in range(self.reshape.target_shape[0])]
        image_emb = tf.keras.backend.stack(image_embeds, axis=1)
        
        # handle text
        text_emb = tf.gather(self.text_embedding, text)
        
        # concat and feed to bert
        concat_emb = self.concat([text_emb, image_emb])
        seg_ids = np.concatenate((np.zeros(max_len, dtype=np.int64), 
                                  np.ones(self.reshape.target_shape[0], dtype=np.int64)))
        print('input shapes to xlnet', concat_emb.shape, seg_ids.shape)
        bert_encodings = self.bert(inputs={'inputs_embeds': concat_emb,
                                            'token_type_ids': seg_ids})[0]
        doc_encoding = tf.squeeze(bert_encodings[:, 0:1, :], axis=1)
        doc_encoding = self.dropout(doc_encoding)
        
        output = self.denseout(doc_encoding)
        
        return output
`

In the line that prints "input shapes to xlnet", I get (None, 116, 768) for the inputs_embeds and (116,) for the token_type_ids, which I expect because I have 100 word embeddings and 16 image embeddings. When I call fit() on this model, it gives the error:


> ValueError: in converted code:
> 
>     <ipython-input-12-b6bccd3c83e0>:50 call  *
>         bert_encodings = self.bert(inputs={'inputs_embeds': concat_emb,
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py:891 __call__
>         outputs = self.call(cast_inputs, *args, **kwargs)
>     /homes/awl27/python36env/lib/python3.6/site-packages/transformers/modeling_tf_xlnet.py:824 call  *
>         outputs = self.transformer(inputs, **kwargs)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/keras/engine/base_layer.py:842 __call__
>         outputs = call_fn(cast_inputs, *args, **kwargs)
>     /homes/awl27/python36env/lib/python3.6/site-packages/transformers/modeling_tf_xlnet.py:530 call  *
>         token_type_ids = tf.transpose(token_type_ids, perm=(1, 0)) if token_type_ids is not None else None
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/ops/array_ops.py:1780 transpose_v2
>         return transpose(a=a, perm=perm, name=name, conjugate=conjugate)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/ops/array_ops.py:1870 transpose
>         ret = transpose_fn(a, perm, name=name)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/ops/gen_array_ops.py:11455 transpose
>         "Transpose", x=x, perm=perm, name=name)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/framework/op_def_library.py:793 _apply_op_helper
>         op_def=op_def)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/framework/func_graph.py:548 create_op
>         compute_device)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:3429 _create_op_internal
>         op_def=op_def)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:1773 __init__
>         control_input_ops)
>     /homes/awl27/python36env/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:1613 _create_c_op
>         raise ValueError(str(e))
> 
>     ValueError: Dimension must be 1 but is 2 for 'transformer/transpose_1' (op: 'Transpose') with input shapes: [116], [2].
> 

## Expected behavior

I expected this to work just like TFBertModel did. If I just change the definition in the BERT class to use TFBertModel instead of TFXLNetModel, it works fine.

## Environment info
PyTorch version: 1.5.0
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 16.04.6 LTS
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.11) 5.4.0 20160609
CMake version: version 3.5.1

Python version: 3.6
Is CUDA available: No
CUDA runtime version: 10.0.130
GPU models and configuration: GPU 0: Tesla K40c
Nvidia driver version: 418.87.00
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.3

Versions of relevant libraries:
[pip3] numpy==1.16.6
[pip3] torch==1.5.0
[pip3] torchtext==0.5.0
[pip3] torchvision==0.6.0
[conda] Could not collect


Comments:
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.