trajdl.algorithms.loc_pred.stlstm module

class trajdl.algorithms.loc_pred.stlstm.HSTLSTM(tokenizer: AbstractTokenizer, embedding_dim: int, hidden_size: int, ts_buckets: Bucketizer, loc_buckets: Bucketizer, optimizer_type: str = 'adam', learning_rate: float = 0.001)

Bases: BaseLightningModel

forward(loc_sessions: List[LongTensor], ts_upper_sessions: List[LongTensor], ts_lower_sessions: List[LongTensor], sd_upper_sessions: List[LongTensor], sd_lower_sessions: List[LongTensor], valid_lengths: List[List[int]])

Same as torch.nn.Module.forward().

Parameters:

• *args – Whatever you decide to pass into the forward method.

• **kwargs – Keyword arguments are also possible.

Returns:

Your model’s output

training_step(batch, batch_idx: int)

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Parameters:

• batch – The output of your data iterable, normally a DataLoader.

• batch_idx – The index of this batch.

• dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Returns:

• Tensor - The loss tensor

• dict - A dictionary which can include any keys, but must include the key 'loss' in the case of automatic optimization.

• None - In automatic optimization, this will skip to the next batch (but is not supported for multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning the loss is not required.

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx):
    x, y, z = batch
    out = self.encoder(x)
    loss = self.loss(out, x)
    return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def __init__(self):
    super().__init__()
    self.automatic_optimization = False


# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx):
    opt1, opt2 = self.optimizers()

    # do training_step with encoder
    ...
    opt1.step()
    # do training_step with decoder
    ...
    opt2.step()

Note

When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.

class trajdl.algorithms.loc_pred.stlstm.STLSTM(tokenizer: AbstractTokenizer, embedding_dim: int, hidden_size: int, ts_bucketizer: Bucketizer, loc_bucketizer: Bucketizer, loc_emb_layer: BaseTokenEmbeddingLayer | None = None)

Bases: Module

cell_step(loc: LongTensor, td_upper: LongTensor, td_lower: LongTensor, sd_upper: LongTensor, sd_lower: LongTensor, hidden: Tuple[Tensor, Tensor], first_step: bool = False)

loc, shape is (B,)，位置batch td_upper, td_lower: shape is (B,)，这个是两个时间步之间的时间差被上界减去和减去下界的batch sd_upper, sd_lower: shape is (B,)，这个是两个时间步之间的空间位移被上界减去和减去下界的batch hidden: [hidden, cell], shape: (B, H), (B, H) first_step: bool, optional

default False

forward(sample: STLSTMSample, hidden: Tuple[Tensor, Tensor] | None = None)

Parameters:

• sample (STLSTMSample) – 输入序列

• hidden (Tuple[torch.Tensor, torch.Tensor], optional) – (hidden, cell), shape: (B, H), (B, H)

property hidden_size: int

init_hidden(batch_size: int, device: device)

property sd_lower: float

property sd_upper: float

property td_lower: float

property td_upper: float

class trajdl.algorithms.loc_pred.stlstm.STLSTMModule(tokenizer: str | AbstractTokenizer, embedding_dim: int, hidden_size: int, ts_bucketizer: str | Bucketizer, loc_bucketizer: str | Bucketizer, reduction: str | LossEnum = 'mean', use_sampled_softmax: bool = True, num_neg_samples: int = 64, loc_emb_layer: BaseTokenEmbeddingLayer | None = None, optimizer_type: str = 'adam', learning_rate: float = 0.001)

Bases: BaseLightningModel

compute_loss(sample: STLSTMSample)

计算一个batch的loss

Parameters:

sample (STLSTMSample) – 需要计算损失的样本

encode(sample: STLSTMSample)

以0为隐藏状态和细胞状态作为初始化，将一个session的前N-1个位置作为输入，输出后N-1个位置的隐藏状态

Parameters:

loc_seq (torch.LongTensor) – 位置序列

Returns:

output – shape is (B, T, H)

Return type:

torch.Tensor

forward(sample: STLSTMSample, k: int = 1)

推理，选择最后k个时间步的输出

Parameters:

• sample (STLSTMSample) – 推理时的样本

• k (int, optional) – 取每条序列的最后k个时间步进行输出，默认值是1

Return type:

torch.Tensor, shape is (B, num_locs)

property hidden_size: int

property loss_reduction: LossEnum

property num_locs: int

on_test_epoch_end()

Called in the test loop at the very end of the epoch.

on_test_epoch_start()

Called in the test loop at the very beginning of the epoch.

on_validation_epoch_end()

Called in the validation loop at the very end of the epoch.

on_validation_epoch_start()

Called in the validation loop at the very beginning of the epoch.

test_step(batch: STLSTMSample, batch_idx: int) → None

Operates on a single batch of data from the test set. In this step you’d normally generate examples or calculate anything of interest such as accuracy.

Parameters:

• batch – The output of your data iterable, normally a DataLoader.

• batch_idx – The index of this batch.

• dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Returns:

• Tensor - The loss tensor

• dict - A dictionary. Can include any keys, but must include the key 'loss'.

• None - Skip to the next batch.

# if you have one test dataloader:
def test_step(self, batch, batch_idx): ...


# if you have multiple test dataloaders:
def test_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single test dataset
def test_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'test_loss': loss, 'test_acc': test_acc})

If you pass in multiple test dataloaders, test_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple test dataloaders
def test_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

Note

If you don’t need to test you don’t need to implement this method.

Note

When the test_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of the test epoch, the model goes back to training mode and gradients are enabled.

training_step(batch: STLSTMSample, batch_idx: int)

Here you compute and return the training loss and some additional metrics for e.g. the progress bar or logger.

Parameters:

• batch – The output of your data iterable, normally a DataLoader.

• batch_idx – The index of this batch.

• dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Returns:

• Tensor - The loss tensor

• dict - A dictionary which can include any keys, but must include the key 'loss' in the case of automatic optimization.

• None - In automatic optimization, this will skip to the next batch (but is not supported for multi-GPU, TPU, or DeepSpeed). For manual optimization, this has no special meaning, as returning the loss is not required.

In this step you’d normally do the forward pass and calculate the loss for a batch. You can also do fancier things like multiple forward passes or something model specific.

Example:

def training_step(self, batch, batch_idx):
    x, y, z = batch
    out = self.encoder(x)
    loss = self.loss(out, x)
    return loss

To use multiple optimizers, you can switch to ‘manual optimization’ and control their stepping:

def __init__(self):
    super().__init__()
    self.automatic_optimization = False


# Multiple optimizers (e.g.: GANs)
def training_step(self, batch, batch_idx):
    opt1, opt2 = self.optimizers()

    # do training_step with encoder
    ...
    opt1.step()
    # do training_step with decoder
    ...
    opt2.step()

Note

When accumulate_grad_batches > 1, the loss returned here will be automatically normalized by accumulate_grad_batches internally.

property use_sampled_softmax: bool

validation_step(batch: STLSTMSample, batch_idx: int) → Tensor

Operates on a single batch of data from the validation set. In this step you’d might generate examples or calculate anything of interest like accuracy.

Parameters:

• batch – The output of your data iterable, normally a DataLoader.

• batch_idx – The index of this batch.

• dataloader_idx – The index of the dataloader that produced this batch. (only if multiple dataloaders used)

Returns:

• Tensor - The loss tensor

• dict - A dictionary. Can include any keys, but must include the key 'loss'.

• None - Skip to the next batch.

# if you have one val dataloader:
def validation_step(self, batch, batch_idx): ...


# if you have multiple val dataloaders:
def validation_step(self, batch, batch_idx, dataloader_idx=0): ...

Examples:

# CASE 1: A single validation dataset
def validation_step(self, batch, batch_idx):
    x, y = batch

    # implement your own
    out = self(x)
    loss = self.loss(out, y)

    # log 6 example images
    # or generated text... or whatever
    sample_imgs = x[:6]
    grid = torchvision.utils.make_grid(sample_imgs)
    self.logger.experiment.add_image('example_images', grid, 0)

    # calculate acc
    labels_hat = torch.argmax(out, dim=1)
    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)

    # log the outputs!
    self.log_dict({'val_loss': loss, 'val_acc': val_acc})

If you pass in multiple val dataloaders, validation_step() will have an additional argument. We recommend setting the default value of 0 so that you can quickly switch between single and multiple dataloaders.

# CASE 2: multiple validation dataloaders
def validation_step(self, batch, batch_idx, dataloader_idx=0):
    # dataloader_idx tells you which dataset this is.
    ...

Note

If you don’t need to validate you don’t need to implement this method.

Note

When the validation_step() is called, the model has been put in eval mode and PyTorch gradients have been disabled. At the end of validation, the model goes back to training mode and gradients are enabled.