| 62 | |
| 63 | |
| 64 | class BertEncoderLayer(Module): |
| 65 | |
| 66 | def __init__(self, |
| 67 | hidden_size, |
| 68 | num_attention_heads, |
| 69 | max_position_embeddings, |
| 70 | hidden_act='relu', |
| 71 | tp_group=None, |
| 72 | tp_size=1, |
| 73 | dtype=None): |
| 74 | super().__init__() |
| 75 | self.input_layernorm = LayerNorm(normalized_shape=hidden_size, |
| 76 | dtype=dtype) |
| 77 | |
| 78 | self.attention = BertAttention( |
| 79 | hidden_size=hidden_size, |
| 80 | num_attention_heads=num_attention_heads, |
| 81 | max_position_embeddings=max_position_embeddings, |
| 82 | tp_group=tp_group, |
| 83 | tp_size=tp_size, |
| 84 | dtype=dtype) |
| 85 | self.mlp = MLP(hidden_size=hidden_size, |
| 86 | ffn_hidden_size=hidden_size * 4, |
| 87 | hidden_act=hidden_act, |
| 88 | tp_group=tp_group, |
| 89 | tp_size=tp_size, |
| 90 | dtype=dtype) |
| 91 | self.post_layernorm = LayerNorm(normalized_shape=hidden_size, |
| 92 | dtype=dtype) |
| 93 | |
| 94 | def forward(self, |
| 95 | hidden_states, |
| 96 | attention_mask=None, |
| 97 | input_lengths=None, |
| 98 | max_input_length=None): |
| 99 | residual = hidden_states |
| 100 | |
| 101 | attention_output = self.attention(hidden_states, |
| 102 | attention_mask=attention_mask, |
| 103 | input_lengths=input_lengths, |
| 104 | max_input_length=max_input_length) |
| 105 | |
| 106 | hidden_states = residual + attention_output |
| 107 | |
| 108 | hidden_states = self.input_layernorm(hidden_states) |
| 109 | |
| 110 | residual = hidden_states |
| 111 | |
| 112 | hidden_states = self.mlp(hidden_states) |
| 113 | |
| 114 | hidden_states = residual + hidden_states |
| 115 | |
| 116 | hidden_states = self.post_layernorm(hidden_states) |
| 117 | |
| 118 | return hidden_states |
| 119 | |
| 120 | |
| 121 | class BertBase(PretrainedModel): |