MultiEmbedding#

class pytorch_forecasting.models.nn.embeddings.MultiEmbedding(embedding_sizes: Dict[str, Tuple[int, int]] | Dict[str, int] | List[int] | List[Tuple[int, int]], x_categoricals: List[str] = None, categorical_groups: Dict[str, List[str]] | None = None, embedding_paddings: List[str] | None = None, max_embedding_size: int = None)[source]#

Bases: Module

Embedding layer for categorical variables including groups of categorical variables.

Enabled for static and dynamic categories (i.e. 3 dimensions for batch x time x categories).

Parameters:

  • embedding_sizes (Union[Dict[str, Tuple[int, int]], Dict[str, int], List[int], List[Tuple[int, int]]]) –

    either

    • dictionary of embedding sizes, e.g. {'cat1': (10, 3)} indicates that the first categorical variable has 10 unique values which are mapped to 3 embedding dimensions. Use get_embedding_size() to automatically obtain reasonable embedding sizes depending on the number of categories.

    • dictionary of categorical sizes, e.g. {'cat1': 10} where embedding sizes are inferred by get_embedding_size().

    • list of embedding and categorical sizes, e.g. [(10, 3), (20, 2)] (requires x_categoricals to be empty)

    • list of categorical sizes where embedding sizes are inferred by get_embedding_size() (requires x_categoricals to be empty).

    If input is provided as list, output will be a single tensor of shape batch x (optional) time x sum(embedding_sizes). Otherwise, output is a dictionary of embedding tensors.

  • x_categoricals (List[str]) – list of categorical variables that are used as input.

  • categorical_groups (Dict[str, List[str]]) – dictionary of categories that should be summed up in an embedding bag, e.g. {'cat1': ['cat2', 'cat3']} indicates that a new categorical variable 'cat1' is mapped to an embedding bag containing the second and third categorical variables. Defaults to empty dictionary.

  • embedding_paddings (List[str]) – list of categorical variables for which the value 0 is mapped to a zero embedding vector. Defaults to empty list.

  • max_embedding_size (int, optional) – if embedding size defined by embedding_sizes is larger than max_embedding_size, it will be constrained. Defaults to None.

Methods

forward(x)

init_embeddings()

items()

keys()

names()

values()

forward(x: Tensor) → Dict[str, Tensor][source]#
Parameters:

x (torch.Tensor) – input tensor of shape batch x (optional) time x categoricals in the order of x_categoricals.

Returns:

dictionary of category names to embeddings

of shape batch x (optional) time x embedding_size if embedding_size is given as dictionary. Otherwise, returns the embedding of shape batch x (optional) time x sum(embedding_sizes). Query attribute output_size to get the size of the output(s).

Return type:

Union[Dict[str, torch.Tensor], torch.Tensor]