huggingface / huggingface/transformers

Modify BERT encoder layers?

Open
#12,241 1 comment 0 reactions 0 assignees View on GitHub
New model
Dominant language
Python
Stars
166k
Forks
34.6k
Avg merge
3d 8h
Merged PRs (30d)
276

Description

Hello I would like to modify the encoder layers of the BERT model, to insert FC and ReLu layers.

This idea allows you to reproduce the use of [Squeeze-and-Excitation Networks](https://arxiv.org/abs/1709.01507)

How to use an nn.module class to handle encoder outputs?

Example:

```
import torch.nn as nn
from transformers import BertModel

class CustomBERTModel(nn.Module):
def __init__(self):
super(CustomBERTModel, self).__init__()
self.bert = BertModel.from_pretrained("bert-base-uncased")
# add your additional layers here, for example a dropout layer followed by a linear classification head
self.dropout = nn.Dropout(0.3)
self.out = nn.Linear(768, 2)

def forward(self, ids, mask, token_type_ids):
sequence_output, pooled_output = self.bert(
ids,
attention_mask=mask,
token_type_ids=token_type_ids
)

# we apply dropout to the sequence output, tensor has shape (batch_size, sequence_length, 768)
sequence_output = self.dropout(sequence_output)

# next, we apply the linear layer. The linear layer (which applies a linear transformation)
# takes as input the hidden states of all tokens (so seq_len times a vector of size 768, each corresponding to
# a single token in the input sequence) and outputs 2 numbers (scores, or logits) for every token
# so the logits are of shape (batch_size, sequence_length, 2)
logits = self.out(sequence_output)

return logits
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.