Add Chronos2Pipeline.embed #361
Conversation
aedade2 to
2b0bf77
Compare
9cc8d98 to
b0f9622
Compare
| def encode( | ||
| self, | ||
| context: torch.Tensor, | ||
| context_mask: torch.Tensor | None = None, | ||
| group_ids: torch.Tensor | None = None, | ||
| future_covariates: torch.Tensor | None = None, | ||
| future_covariates_mask: torch.Tensor | None = None, | ||
| num_output_patches: int = 1, | ||
| future_target: torch.Tensor | None = None, | ||
| future_target_mask: torch.Tensor | None = None, | ||
| output_attentions: bool = False, | ||
| ): |
There was a problem hiding this comment.
I wish the diff would be more helpful here: is the body of this simply moved from forward?
There was a problem hiding this comment.
Yes, the first (encoding) portion from forward has been factored out into encode.
| or 2-dimensional of shape (n_variates, history_length). The history_lengths may be different across elements; left-padding | ||
| will be applied, if needed. | ||
| batch_size | ||
| The batch size used for generating embeddings. Note that the batch size here means the total number of time series which are input into the model. |
There was a problem hiding this comment.
I'm not sure this is clear to me: does the batch_size refer to the .shape[0] of the tensors being processed? Or does it span the variates dimension as well? (.shape[1]) I suppose it's the latter, given the docstring for the dataset class:
chronos-forecasting/src/chronos/chronos2/dataset.py
Lines 402 to 405 in e48f480
I see this is pretty much the description of batch_size everywhere (here, predict methods, dataset class). Maybe the confusion comes from "total number of time series" instead of "total number of variates", or something like that. But this could also be addressed separately.
There was a problem hiding this comment.
Internally, there's no notion of a variate dimension in the model: only batch and time (patch) axes. The batch_size here refers to the maximum items x (co)-variates per batch. Open to suggestions on a better docstring.
abdulfatir
left a comment
There was a problem hiding this comment.
Thanks @lostella.
| def encode( | ||
| self, | ||
| context: torch.Tensor, | ||
| context_mask: torch.Tensor | None = None, | ||
| group_ids: torch.Tensor | None = None, | ||
| future_covariates: torch.Tensor | None = None, | ||
| future_covariates_mask: torch.Tensor | None = None, | ||
| num_output_patches: int = 1, | ||
| future_target: torch.Tensor | None = None, | ||
| future_target_mask: torch.Tensor | None = None, | ||
| output_attentions: bool = False, | ||
| ): |
There was a problem hiding this comment.
Yes, the first (encoding) portion from forward has been factored out into encode.
| or 2-dimensional of shape (n_variates, history_length). The history_lengths may be different across elements; left-padding | ||
| will be applied, if needed. | ||
| batch_size | ||
| The batch size used for generating embeddings. Note that the batch size here means the total number of time series which are input into the model. |
There was a problem hiding this comment.
Internally, there's no notion of a variate dimension in the model: only batch and time (patch) axes. The batch_size here refers to the maximum items x (co)-variates per batch. Open to suggestions on a better docstring.
Issue #, if available: #354
Description of changes: This PR adds
Chronos2Pipeline.embedto enable users to extract embeddings from the last encoder layer in an easy way. The API and behavior is similar to what Chronos and Chronos-Bolt provides.By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.