trajdl.algorithms.loc_pred.rnn module

class trajdl.algorithms.loc_pred.rnn.RNNNextLocPredictor(embedding_layer: Embedding, rnn_hidden_size: int, fc_hidden_size: int, output_size: int, num_layers: int, padding_value: int, dropout: float = 0.0)

Bases: Module

forward(src: LongTensor, lengths: List[int])

Parameters:

• src (shape is (B, T))

• lengths (List[int])

trajdl.algorithms.loc_pred.rnn.select_last_k_elements(x: Tensor, lengths: List[int], k: int) → Tensor

选出序列里面最后k个元素，这里要考虑每个序列的实际长度，x是padding后的序列 shape of x is (B, T, *)