CODEWITHSUNDEEP
tensorflowkerasnlp
Alexey Kvashchuk
Alexey Kvashchuk

Reputation: 123

How to access a particular layer of Huggingface's pre-trained BERT model?

For experimentation purposes, I need to access an Embedding layer of the encoder. That is, assuming Tensorflow implementation, the layer defined as tf.keras.layers.Embedding(...).

For example, what is a way to set 'embeddings_regularizer=' argument of the Embedding() layer in the encoder part of the transformer?

Upvotes: 2

Views: 6986

Answers (1)

ML_Engine
ML_Engine

Reputation: 1185

You can iterate over the BERT model in the same way as any other model, like so:

for layer in model.layers:
    if isinstance(layer ,tf.keras.layers.Embedding):
        layer.embeddings_regularizer = argument

isinstance checks the type of the layer, so really you can put any layer type here and change what you need.

I haven't checked specifically whether embeddings_regularizer is available, however if you want to see what methods are available to that particular layer, run a debugger and call dir(layer) inside the above function.

Updated question

The TFBertForSequenceClassification model has 3 layers:

>>> model.summary()

Model: "tf_bert_for_sequence_classification"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
bert (TFBertMainLayer)       multiple                  108310272 
_________________________________________________________________
dropout_37 (Dropout)         multiple                  0         
_________________________________________________________________
classifier (Dense)           multiple                  1538      
=================================================================
Total params: 108,311,810
Trainable params: 108,311,810
Non-trainable params: 0

Similarly, calling model.layers gives:

[<transformers.models.bert.modeling_tf_bert.TFBertMainLayer at 0x7efda85595d0>,
 <tensorflow.python.keras.layers.core.Dropout at 0x7efd6000ae10>,
 <tensorflow.python.keras.layers.core.Dense at 0x7efd6000afd0>]

We can access the layers inside TFBertMainLayer:

>>> model.layers[0]._layers


[<transformers.models.bert.modeling_tf_bert.TFBertEmbeddings at 0x7efda8080f90>,
 <transformers.models.bert.modeling_tf_bert.TFBertEncoder at 0x7efda855ced0>,
 <transformers.models.bert.modeling_tf_bert.TFBertPooler at 0x7efda84f0450>,
 DictWrapper({'name': 'bert'})]

So from the above we can access the TFBertEmbeddings layer by:

model.layers[0].embeddings

OR

model.layers[0]._layers[0]

If you check the documentation (search for the "TFBertEmbeddings" class) you can see that this inherits a standard tf.keras.layers.Layer which means you have access to all the normal regularizer methods, so you should be able to call something like:

from tensorflow.keras import regularizers

model.layers[0].embeddings.activity_regularizer = regularizers.l2(1e-5)

Or whatever argument / regularizer you need to change. See here for regularizer docs.

Upvotes: 5

Related Questions