attention-is-all-you-need-pytorch 源码阅读

训练数据流

train.train_epoch

对training_data进行迭代, 产生batch, 其中有src_seq, trg_seq

src_seq.shape
Out[11]: torch.Size([256, 32])
src_seq
Out[12]: 
tensor([[   2, 4567, 4578,  ...,    1,    1,    1],
        [   2, 4558, 4565,  ...,    1,    1,    1],
        [   2, 4558, 4565,  ...,    1,    1,    1],
        ...,
        [   2, 4558,   64,  ...,    1,    1,    1],
        [   2, 4564, 5051,  ...,    1,    1,    1],
        [   2, 4567, 4578,  ...,    1,    1,    1]])

2是开始, 1是结束, 32是句子长度, 256是batch数

Transformer

transformer.Models.Transformer.forward

现在数据被丢进了Transformer这个模型

src_mask = get_pad_mask(src_seq, self.src_pad_idx)
trg_mask = get_pad_mask(trg_seq, self.trg_pad_idx) & get_subsequent_mask(trg_seq)

enc_output, *_ = self.encoder(src_seq, src_mask)

src_mask会在ScaledDotProductAttention用到

编码了一波, enc_output其实与输入数据的size一样

enc_output.shape
Out[9]: torch.Size([256, 36, 512])

去看Decoder

Encoder

transformer.Models.Encoder

#                               9521         512          1
self.src_word_emb = nn.Embedding(n_src_vocab, d_word_vec, padding_idx=pad_idx)

what-does-padding-idx-do-in-nn-embeddings

transformer.Models.Encoder.forward

enc_output = self.dropout(self.position_enc(self.src_word_emb(src_seq)))

512维词向量

enc_output.shape
Out[16]: torch.Size([