TransformerEncoder is a stack of N encoder layers Args: src_n_token: the source vocab size d_model: the number of expected features in the encoder inputs (default=512). n_head: the number of heads in the multi head attention models (default=8). dim
| 253 | |
| 254 | |
| 255 | class TransformerEncoder(layer.Layer): |
| 256 | """TransformerEncoder is a stack of N encoder layers |
| 257 | Args: |
| 258 | src_n_token: the source vocab size |
| 259 | d_model: the number of expected features in the encoder inputs (default=512). |
| 260 | n_head: the number of heads in the multi head attention models (default=8). |
| 261 | dim_feedforward: the dimension of the feedforward network model (default=2048). |
| 262 | n_layers: the number of sub-encoder-layers in the encoder (default=6). |
| 263 | """ |
| 264 | |
| 265 | def __init__(self, src_n_token, d_model=512, n_head=8, dim_feedforward=2048, n_layers=6): |
| 266 | super(TransformerEncoder, self).__init__() |
| 267 | self.src_n_token = src_n_token |
| 268 | self.d_model = d_model |
| 269 | self.n_head = n_head |
| 270 | self.dim_feedforward = dim_feedforward |
| 271 | self.n_layers = n_layers |
| 272 | |
| 273 | # input_emb / pos_emb / n-encoder layers |
| 274 | self.input_emb = layer.Embedding(input_dim=src_n_token, output_dim=d_model) |
| 275 | self.pos_emb = layer.Embedding(input_dim=src_n_token, output_dim=d_model) |
| 276 | self.layers = [] |
| 277 | for _ in range(self.n_layers): |
| 278 | self.layers.append(TransformerEncoderLayer(d_model=d_model, n_head=n_head, dim_feedforward=dim_feedforward)) |
| 279 | |
| 280 | def forward(self, enc_inputs): |
| 281 | """Pass the input through the encoder in turn. |
| 282 | Args: |
| 283 | enc_inputs: the sequence to the encoder (required). [batch_size, src_len] |
| 284 | """ |
| 285 | # [batch_size, src_len, d_model] |
| 286 | word_emb = self.input_emb(enc_inputs) |
| 287 | |
| 288 | self.pos_emb.initialize(enc_inputs) |
| 289 | self.pos_emb.from_pretrained(W=TransformerEncoder._get_sinusoid_encoding_table(self.src_n_token, self.d_model), freeze=True) |
| 290 | # [batch_size, src_len, d_model] |
| 291 | pos_emb = self.pos_emb(enc_inputs) |
| 292 | # enc_outputs [batch_size, src_len, d_model] |
| 293 | enc_outputs = autograd.add(word_emb, pos_emb) |
| 294 | |
| 295 | # enc_self_attn_mask [batch_size, src_len, src_len] |
| 296 | enc_self_attn_mask = TransformerEncoder._get_attn_pad_mask(enc_inputs, enc_inputs) |
| 297 | |
| 298 | enc_self_attns = [] |
| 299 | for layer in self.layers: |
| 300 | enc_outputs, enc_self_attn = layer(enc_outputs, enc_self_attn_mask) |
| 301 | enc_self_attns.append(enc_self_attn) |
| 302 | return enc_outputs, enc_self_attns |
| 303 | |
| 304 | @staticmethod |
| 305 | def _get_attn_pad_mask(seq_q, seq_k): |
| 306 | """ |
| 307 | Args: |
| 308 | seq_q: [batch_size, seq_len] |
| 309 | seq_k: [batch_size, seq_len] |
| 310 | Returns: [batch_size, seq_len, seq_len] |
| 311 | """ |
| 312 | batch_size, len_q = seq_q.shape |