| 349 | |
| 350 | class BertSelfAttention(nn.Module): |
| 351 | def __init__(self, config): |
| 352 | super(BertSelfAttention, self).__init__() |
| 353 | if config.hidden_size % config.num_attention_heads != 0: |
| 354 | raise ValueError( |
| 355 | "The hidden size (%d) is not a multiple of the number of attention " |
| 356 | "heads (%d)" % (config.hidden_size, config.num_attention_heads)) |
| 357 | self.num_attention_heads = config.num_attention_heads |
| 358 | self.attention_head_size = int(config.hidden_size / config.num_attention_heads) |
| 359 | self.all_head_size = self.num_attention_heads * self.attention_head_size |
| 360 | |
| 361 | self.query = nn.Linear(config.hidden_size, self.all_head_size) |
| 362 | self.key = nn.Linear(config.hidden_size, self.all_head_size) |
| 363 | self.value = nn.Linear(config.hidden_size, self.all_head_size) |
| 364 | |
| 365 | self.dropout = nn.Dropout(config.attention_probs_dropout_prob) |
| 366 | |
| 367 | def transpose_for_scores(self, x): |
| 368 | new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size) |