(
self,
hidden_size,
num_attention_heads,
layer_number,
layernorm_epsilon=1e-5,
fp16=True,
attention_softmax_in_fp32=True,
)
| 406 | """ |
| 407 | |
| 408 | def __init__( |
| 409 | self, |
| 410 | hidden_size, |
| 411 | num_attention_heads, |
| 412 | layer_number, |
| 413 | layernorm_epsilon=1e-5, |
| 414 | fp16=True, |
| 415 | attention_softmax_in_fp32=True, |
| 416 | ): |
| 417 | super(TransformerLayer, self).__init__() |
| 418 | self.hidden_size = hidden_size |
| 419 | self.layernorm_epsilon = layernorm_epsilon |
| 420 | self.layer_number = layer_number |
| 421 | |
| 422 | # Layernorm on the input data. |
| 423 | self.input_layernorm = torch.nn.LayerNorm(hidden_size, |
| 424 | eps=self.layernorm_epsilon) |
| 425 | |
| 426 | # Self attention. |
| 427 | self.attention = SelfAttention(hidden_size, |
| 428 | num_attention_heads, |
| 429 | layer_number, |
| 430 | fp16, |
| 431 | attention_softmax_in_fp32) |
| 432 | |
| 433 | # Layernorm on the input data. |
| 434 | self.post_attention_layernorm = torch.nn.LayerNorm(self.hidden_size, |
| 435 | eps=self.layernorm_epsilon) |
| 436 | self.mlp = MLP(self.hidden_size) |
| 437 | |
| 438 | def forward( |
| 439 | self, |
nothing calls this directly
no test coverage detected