RNN Keras model for NLP takes lot of time while training with no decrease in validation loss

user_8703627 0
Viewed 331

I have built an RNN model for entity recognition. I used BERT embedding and then processed the results through a RNN Model. However, while training the model for 5 epochs, each epoch seems to take about 2 hours. Also, the validation loss does not seem to decrease at all.

I am running the process on a RTX 2080 GPU. I have tried manipulating the model but does not improve the model. The dataset I have is of about 400000 sentences.

This is my model:

def build_model(max_seq_length, n_tags): 
    in_id = Input(shape=(max_seq_length,), name="input_ids")
    in_mask = Input(shape=(max_seq_length,), name="input_masks")
    in_segment = Input(shape=(max_seq_length,), name="segment_ids")

    bert_inputs = [in_id, in_mask, in_segment]   
    bert_output = BertLayer(n_fine_tune_layers=3, pooling="first")(bert_inputs)
    x = RepeatVector(max_seq_length)(bert_output)
    x = Bidirectional(LSTM(units=lstm_units, return_sequences=True,
                           recurrent_dropout=0.2, dropout=0.2))(x)
    x_rnn = Bidirectional(LSTM(units=lstm_units, return_sequences=True,
                               recurrent_dropout=0.2, dropout=0.2))(x)
    x = add([x, x_rnn])  # residual connection to the first biLSTM
    pred = TimeDistributed(Dense(n_tags, activation="softmax"))(x)

    model = Model(inputs=bert_inputs, outputs=pred)
    model.compile(loss='sparse_categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
    model.summary()
    return model

This is the model summary:

Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_ids (InputLayer)          (None, 30)           0                                            
__________________________________________________________________________________________________
input_masks (InputLayer)        (None, 30)           0                                            
__________________________________________________________________________________________________
segment_ids (InputLayer)        (None, 30)           0                                            
__________________________________________________________________________________________________
bert_layer_3 (BertLayer)        ((None, 30), 768)    110104890   input_ids[0][0]                  
                                                                 input_masks[0][0]                
                                                                 segment_ids[0][0]                
__________________________________________________________________________________________________
repeat_vector_2 (RepeatVector)  ((None, 30), 30, 768 0           bert_layer_3[0][0]               
__________________________________________________________________________________________________
bidirectional_2 (Bidirectional) ((None, 30), 30, 200 695200      repeat_vector_2[0][0]            
__________________________________________________________________________________________________
bidirectional_3 (Bidirectional) ((None, 30), 30, 200 240800      bidirectional_2[0][0]            
__________________________________________________________________________________________________
add_1 (Add)                     ((None, 30), 30, 200 0           bidirectional_2[0][0]            
                                                                 bidirectional_3[0][0]            
__________________________________________________________________________________________________
time_distributed_1 (TimeDistrib ((None, 30), 30, 3)  603         add_1[0][0]                      
==================================================================================================
Total params: 111,041,493
Trainable params: 22,790,811
Non-trainable params: 88,250,682
__________________________________________________________________________________________________

Logs:

 32336/445607 [=>............................] - ETA: 2:12:59 - loss: 0.3469 - acc: 0.9068
 32352/445607 [=>............................] - ETA: 2:12:58 - loss: 0.3469 - acc: 0.9068
 32368/445607 [=>............................] - ETA: 2:12:58 - loss: 0.3469 - acc: 0.9068

Can you help me find out where I am going wrong?

recurrent-neural-network lstm keras deep-learning nlp
1 Answers

If you are using Bert for embedding, the output shape should be (None, 30, 768). But your Bert model is returning a (None, 768) tensor, then you used RepeatVector to duplicate it. I guess you are extracting [CLS] output from Bert. Please extract the right layer from Bert model.

And the reason why the trianing takes so long is just because for every epoch you need to pass all your data through the huge bert model, even you are freezing most of the layers.

user_103719870
Related