| 168 | |
| 169 | |
| 170 | class EncoderLayer(torch.nn.Module): |
| 171 | def __init__(self, d_model_size, num_heads, dff, rate=0.1): |
| 172 | super().__init__() |
| 173 | |
| 174 | self.multi_head_attention = MultiHeadAttention(d_model_size, num_heads) |
| 175 | self.ffn = point_wise_feed_forward_network(d_model_size, dff) |
| 176 | |
| 177 | self.layernorm1 = torch.nn.LayerNorm(d_model_size, eps=1e-6) |
| 178 | self.layernorm2 = torch.nn.LayerNorm(d_model_size, eps=1e-6) |
| 179 | |
| 180 | self.dropout1 = torch.nn.Dropout(rate) |
| 181 | self.dropout2 = torch.nn.Dropout(rate) |
| 182 | |
| 183 | def forward( |
| 184 | self, x, mask, layer_past=None, attention_mask=None, head_mask=None, use_cache=False, output_attentions=False |
| 185 | ): |
| 186 | normed = self.layernorm1(x) |
| 187 | attn_outputs = self.multi_head_attention( |
| 188 | normed, |
| 189 | normed, |
| 190 | normed, |
| 191 | mask, |
| 192 | layer_past=layer_past, |
| 193 | attention_mask=attention_mask, |
| 194 | head_mask=head_mask, |
| 195 | use_cache=use_cache, |
| 196 | output_attentions=output_attentions, |
| 197 | ) |
| 198 | attn_output = attn_outputs[0] |
| 199 | attn_output = self.dropout1(attn_output) |
| 200 | out1 = x + attn_output |
| 201 | |
| 202 | out2 = self.layernorm2(out1) |
| 203 | ffn_output = self.ffn(out2) |
| 204 | ffn_output = self.dropout2(ffn_output) |
| 205 | out2 = out1 + ffn_output |
| 206 | |
| 207 | outputs = (out2,) + attn_outputs[1:] |
| 208 | return outputs |
| 209 | |
| 210 | |
| 211 | class CTRLPreTrainedModel(PreTrainedModel): |