how to extract several layers of BERT or GPT as a new model?
Author: willanxywc
Pull request: no
Labels: name: wontfix

How can I, for example, extract 8 layers from the 12 BertLayers of  the _bert-base-uncased_ to form a new model? I want to use the _embedding_ and _pooler_ layer of orginal model, but use only a portion of the _encoder_ layers.

Comments:
Interesting use-case!

The easiest way would be to simply load to models, one with the `bert-base-cased` checkpoint, the other randomly initialized, and to assign trained layers to the new model. Something like this:

```py
from transformers import BertModel, BertConfig
import torch

bert_base_cased = BertModel.from_pretrained("bert-base-cased")  # Instantiate model using the trained weights
model = BertModel(BertConfig.from_pretrained("bert-base-cased"))  # Randomly initialize model, with the same size as the trained model

layers_to_replace = [1, 2, 3, 8]
for layer in layers_to_replace:
    model.base_model.encoder.layer[layer] = bert_base_cased.base_model.encoder.layer[layer]

# Let's compare the key values of the attention layers to make sure they're the same
i = 0
for original_layer, new_layer in zip(model.base_model.encoder.layer, bert_base_cased.base_model.encoder.layer):
    original_attention_key = original_layer.attention.self.key.weight
    new_attention_key = new_layer.attention.self.key.weight
    difference = (torch.max(torch.abs(original_attention_key - new_attention_key)).item())

    print(f"Layers {i} are {'not ' if difference else ''}the same.")
    i += 1

```

This outputs:

```
Layers 0 are not the same.
Layers 1 are the same.
Layers 2 are the same.
Layers 3 are the same.
Layers 4 are not the same.
Layers 5 are not the same.
Layers 6 are not the same.
Layers 7 are not the same.
Layers 8 are the same.
Layers 9 are not the same.
Layers 10 are not the same.
Layers 11 are not the same.
```

Thanks, but I find this does not work for GPT2LMHeadModel. How could I extract the hidden layers of a GPT2LMHeadModel please?

You can do the same, but the layers are under `model.base_model.h`:

```py
[...]
for layer in layers_to_replace:
    model.base_model.h[layer] = [...]
[...]
```

Thanks @LysandreJik !

I am wondering how you would do this in the keras versions. From tinkering around, I think you access the layers with `model.layers[0].encoder.layer`, since the length of this is 12, so I'm guessing it's for the 12 layers in the Bert model. 

So you would do something like 

```
layers_to_replace = [1, 2, 3, 8]
for layer in layers_to_replace:
    newModel.layers[0].encoder.layer[layer] = trainedModel.layers[0].encoder.layer[layer]
```

Does that seem right to you?

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.


@LysandreJik This solution randomly initialize Embedding layer's weight rather than load from bert pretrained Embedding，which leads to a huge performance decline and confuses me for a week. The correct solution is:
```python
from transformers import BertModel, BertConfig
import torch

bert_version = "bert-base-cased"
bert_base_cased = BertModel.from_pretrained(bert_version)  # Instantiate model using the trained weights
config = BertConfig.from_pretrained(bert_version)
model = BertModel(config=config)  # Randomly initialize model, with the same size as the trained model

# add these two lines
model.embeddings = bert_base_cased.embeddings
model.pooler = bert_base_cased.pooler

layers_to_replace = [1, 2, 3, 8]
for layer in layers_to_replace:
    model.base_model.encoder.layer[layer] = bert_base_cased.base_model.encoder.layer[layer]
```

also，if you just want the first 4 layers, the easier and safer way is:
```python
from transformers import BertModel, BertConfig
import torch

bert_version = "bert-base-cased"
bert_base_cased = BertModel.from_pretrained(bert_version)  # Instantiate model using the trained weights
config = BertConfig.from_pretrained(bert_version)
config.num_hidden_layers = 4
model = BertModel.from_pretrained(bert_version, config=config)  # auto skip unused layers

for param_name in model.state_dict():
    sub_param, full_param = model.state_dict()[param_name], bert_base_cased.state_dict()[param_name] # type: torch.Tensor, torch.Tensor
    assert (sub_param.cpu().numpy() == full_param.cpu().numpy()).all(), param_name
        
```

@dalek-who hey I tried to run your code before my other code to construct the model (and this is all in Sagemaker), however got this error: 

CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

Here's the code I ran after your code:

```
class BERTClass(torch.nn.Module):
    def __init__(self):
        super(BERTClass, self).__init__()
        self.bert_model = model 
        self.dropout = torch.nn.Dropout(0.5)
        self.linear = torch.nn.Linear(768, 9)    
    def forward(self, input_ids, attn_mask, token_type_ids):
        output = self.bert_model(
            input_ids, 
            attention_mask=attn_mask, 
            token_type_ids=token_type_ids
        )
        output_dropout = self.dropout(output.pooler_output)
        output = self.linear(output_dropout)
        return output
bert_model = BERTClass()
bert_model.to(device)
```

Anyone has any idea why?

@Bambry When do you get this error? On construct the model, or on forward?
`CUDA error: device-side assert triggered` often occurs when a layer receives illegal inputs, for example a `BCELoss` receives a illegal label `3`.
Maybe you should check your tensor and parameter's shape, value or dtype.