API

 torch / nn / torch.nn


TransformerEncoder

class torch.nn.TransformerEncoder(encoder_layer, num_layers, norm=None)[source]

TransformerEncoder is a stack of N encoder layers

Parameters
  • encoder_layer – an instance of the TransformerEncoderLayer() class (required).

  • num_layers – the number of sub-encoder-layers in the encoder (required).

  • norm – the layer normalization component (optional).

Examples::
>>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
>>> transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
>>> src = torch.rand(10, 32, 512)
>>> out = transformer_encoder(src)
forward(src: torch.Tensor, mask: Optional[torch.Tensor] = None, src_key_padding_mask: Optional[torch.Tensor] = None) → torch.Tensor[source]

Pass the input through the encoder layers in turn.

Parameters
  • src – the sequence to the encoder (required).

  • mask – the mask for the src sequence (optional).

  • src_key_padding_mask – the mask for the src keys per batch (optional).

Shape:

see the docs in Transformer class.


此页内容是否对您有帮助