l5kit.prediction.vectorized.transformer module¶

class l5kit.prediction.vectorized.transformer.TransformerDecoder(decoder_layer: torch.nn.modules.module.Module, num_layers: int, norm: Optional[torch.nn.modules.module.Module] = None, return_intermediate: Optional[bool] = False)¶

Bases: torch.nn.modules.module.Module

A stack of N transformer decoder layers. See PyTorch docs for more details.

forward(tgt: torch.Tensor, memory: torch.Tensor, tgt_mask: Optional[torch.Tensor] = None, memory_mask: Optional[torch.Tensor] = None, tgt_key_padding_mask: Optional[torch.Tensor] = None, memory_key_padding_mask: Optional[torch.Tensor] = None, pos: Optional[torch.Tensor] = None, query_pos: Optional[torch.Tensor] = None) → torch.Tensor¶

Computes the forward of the module.

Parameters

tgt – decoder layer input, used as input and queries
memory – output of the last encoder layer, used as keys and values
tgt_mask – attention mask for tgt
memory_mask – attention mask for memory
tgt_key_padding_mask – key padding mask for tgt
memory_key_padding_mask – key padding mask for memory
pos – positional embedding tensor, it will be added to memory
query_pos – query positional embedding tensor, it will be added to tgt

Returns

the module output

reset_parameters() → None¶: See https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/transformer.py#L78

training: bool¶

class l5kit.prediction.vectorized.transformer.TransformerDecoderLayer(d_model: int, nhead: int, dim_feedforward: Optional[int] = 2048, dropout: float = 0.1, activation: Optional[str] = 'relu', normalize_before: Optional[bool] = False)¶

Bases: torch.nn.modules.module.Module

Transformer decoder layer composed of self-attn, multi-head-attn and feedforward network. See PyTorch docs for more details.

Computes the forward of the module.

Parameters

tgt – decoder layer input, used as input and queries
memory – output of the last encoder layer, used as keys and values
tgt_mask – attention mask for tgt
memory_mask – attention mask for memory
tgt_key_padding_mask – key padding mask for tgt
memory_key_padding_mask – key padding mask for memory
pos – positional embedding tensor, it will be added to memory
query_pos – query positional embedding tensor, it will be added to tgt

Returns

the layer output

forward_post(tgt: torch.Tensor, memory: torch.Tensor, tgt_mask: Optional[torch.Tensor] = None, memory_mask: Optional[torch.Tensor] = None, tgt_key_padding_mask: Optional[torch.Tensor] = None, memory_key_padding_mask: Optional[torch.Tensor] = None, pos: Optional[torch.Tensor] = None, query_pos: Optional[torch.Tensor] = None) → torch.Tensor¶: Forward normalizing after the self-attention module

forward_pre(tgt: torch.Tensor, memory: torch.Tensor, tgt_mask: Optional[torch.Tensor] = None, memory_mask: Optional[torch.Tensor] = None, tgt_key_padding_mask: Optional[torch.Tensor] = None, memory_key_padding_mask: Optional[torch.Tensor] = None, pos: Optional[torch.Tensor] = None, query_pos: Optional[torch.Tensor] = None) → torch.Tensor¶: Forward normalizing before the self-attention module

training: bool¶

with_pos_embed(tensor: torch.Tensor, pos: Optional[torch.Tensor]) → torch.Tensor¶

class l5kit.prediction.vectorized.transformer.TransformerEncoder(encoder_layer: torch.nn.modules.module.Module, num_layers: int, norm: Optional[torch.nn.modules.module.Module] = None)¶

Bases: torch.nn.modules.module.Module

A stack of N transformer encoder layers. See PyTorch docs for more details.

forward(src: torch.Tensor, mask: Optional[torch.Tensor] = None, src_key_padding_mask: Optional[torch.Tensor] = None, pos: Optional[torch.Tensor] = None) → torch.Tensor¶

Computes the forward of the module.

Parameters

src – module input, shape [num_elements, batch_size, num_features]
mask – attention mask, shape [(batch_size,) num_elements, num_elements]
src_key_padding_mask – key padding mask, shape [batch_size, num_elements]
pos – positional embedding tensor, it will be added to src. shape [num_elements, batch_size, num_features]

Returns

tuple containing the module output and a list of attention weights (one for each encoder layer) shape [num_elements, batch_size, num_features], List[[batch_size, num_elements, num_elements]]

reset_parameters() → None¶

training: bool¶

class l5kit.prediction.vectorized.transformer.TransformerEncoderLayer(d_model: int, nhead: int, dim_feedforward: int = 2048, dropout: float = 0.1, activation: Optional[str] = 'relu', normalize_before: Optional[bool] = False)¶

Bases: torch.nn.modules.module.Module

Transformer encoder layer composed of self-attn and feedforward network. See PyTorch docs for more details.

forward(src: torch.Tensor, src_mask: Optional[torch.Tensor] = None, src_key_padding_mask: Optional[torch.Tensor] = None, pos: Optional[torch.Tensor] = None) → Tuple[torch.Tensor, Optional[torch.Tensor]]¶

Computes the forward of the module.

Parameters

src – input of shape [num_elements, batch_size, num_features]
src_mask – attention mask [batch_size, num_elements, num_elements]
src_key_padding_mask – key padding mask [batch_size, num_elements]
pos – positional embedding tensor, it will be added to src

Returns

tuple containing the layer output and the computed attention weights

forward_post(src: torch.Tensor, src_mask: Optional[torch.Tensor] = None, src_key_padding_mask: Optional[torch.Tensor] = None, pos: Optional[torch.Tensor] = None) → Tuple[torch.Tensor, Optional[torch.Tensor]]¶: Forward normalizing after the self-attention module

forward_pre(src: torch.Tensor, src_mask: Optional[torch.Tensor] = None, src_key_padding_mask: Optional[torch.Tensor] = None, pos: Optional[torch.Tensor] = None) → Tuple[torch.Tensor, Optional[torch.Tensor]]¶: Forward normalizing before the self-attention module

training: bool¶

with_pos_embed(tensor: torch.Tensor, pos: Optional[torch.Tensor]) → torch.Tensor¶