fastNLP.modules.encoder.star_transformer module

undocumented Star-Transformer 的encoder部分的 Pytorch 实现

class fastNLP.modules.encoder.star_transformer.StarTransformer(hidden_size, num_layers, num_head, head_dim, dropout=0.1, max_len=None)[源代码]

基类:torch.nn.modules.module.Module

别名 fastNLP.modules.StarTransformer fastNLP.modules.encoder.StarTransformer

Star-Transformer 的encoder部分。 输入3d的文本输入, 返回相同长度的文本编码

paper: https://arxiv.org/abs/1902.09113

__init__(hidden_size, num_layers, num_head, head_dim, dropout=0.1, max_len=None)[源代码]
参数
  • hidden_size (int) – 输入维度的大小。同时也是输出维度的大小。

  • num_layers (int) – star-transformer的层数

  • num_head (int) – head的数量。

  • head_dim (int) – 每个head的维度大小。

  • dropout (float) – dropout 概率. Default: 0.1

  • max_len (int) – int or None, 如果为int,输入序列的最大长度, 模型会为输入序列加上position embedding。 若为`None`,忽略加上position embedding的步骤. Default: None

forward(data, mask)[源代码]
参数
  • data (FloatTensor) – [batch, length, hidden] 输入的序列

  • mask (ByteTensor) – [batch, length] 输入序列的padding mask, 在没有内容(padding 部分) 为 0, 否则为 1

返回

[batch, length, hidden] 编码后的输出序列

[batch, hidden] 全局 relay 节点, 详见论文

training: bool