| 348 | |
| 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) |
| 369 | x = x.view(*new_x_shape) |
| 370 | return x.permute(0, 2, 1, 3) |
| 371 | |
| 372 | def forward(self, hidden_states, attention_mask): |
| 373 | mixed_query_layer = self.query(hidden_states) |
| 374 | mixed_key_layer = self.key(hidden_states) |
| 375 | mixed_value_layer = self.value(hidden_states) |
| 376 | |
| 377 | query_layer = self.transpose_for_scores(mixed_query_layer) |
| 378 | key_layer = self.transpose_for_scores(mixed_key_layer) |
| 379 | value_layer = self.transpose_for_scores(mixed_value_layer) |
| 380 | |
| 381 | # Take the dot product between "query" and "key" to get the raw attention scores. |
| 382 | attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) |
| 383 | attention_scores = attention_scores / math.sqrt(self.attention_head_size) |
| 384 | # Apply the attention mask is (precomputed for all layers in BertModel forward() function) |
| 385 | attention_scores = attention_scores + attention_mask |
| 386 | |
| 387 | # Normalize the attention scores to probabilities. |
| 388 | attention_probs = nn.Softmax(dim=-1)(attention_scores) |
| 389 | |
| 390 | # This is actually dropping out entire tokens to attend to, which might |
| 391 | # seem a bit unusual, but is taken from the original Transformer paper. |
| 392 | attention_probs = self.dropout(attention_probs) |
| 393 | |
| 394 | previous_type = attention_probs.type() |
| 395 | context_layer = torch.matmul(attention_probs, value_layer) |
| 396 | context_layer = context_layer.permute(0, 2, 1, 3).contiguous() |
| 397 | new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) |
| 398 | context_layer = context_layer.view(*new_context_layer_shape) |
| 399 | return context_layer |
| 400 | |
| 401 | |
| 402 | class BertSelfOutput(nn.Module): |