How to access BERT's inter layer?

Question

I want to put [batch_size, 768, text_length] tensor into 6th layer of BERT.

1. How can I give input to 6th layer?

2. Can I take just 6~last layer of BERT then use it?

Thank you.

Answer 1

If you want to use only the last 6 layers of BERT, you can create a new model that only consists of these layers, using the following code:

import torch
from transformers import BertModel

model = BertModel.from_pretrained("bert-base-uncased")
last_6_layers = torch.nn.ModuleList(model.bert.encoder.layer[-6:]).to(device)

This will create a new model last_6_layers that consists of only the last 6 layers of the BERT model. You can then input your [batch_size, 768, text_length] tensor into this new model and use it for your specific task.