bitorch.layers.bembedding.BEmbeddingBag

class bitorch.layers.bembedding.BEmbeddingBag(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None, weight_quantization: Optional[Union[Quantization, str]] = None, device: Optional[Union[str, device]] = None, sign_bool: bool = False, mode: str = 'mean')

Binarized version of pytorchs embedding bag. Uses given binarization method to binarize the weights. Memory consumption during training increases with batch size. Inference is always small.

Methods

__init__	Initializes internal Module state, shared by both nn.Module and ScriptModule.
apply_aggregate	Aggregates the unique embedding vectors using the defined mode.
forward	Generates embeddings from given tokens and offsets.

Attributes

__init__(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None, weight_quantization: Optional[Union[Quantization, str]] = None, device: Optional[Union[str, device]] = None, sign_bool: bool = False, mode: str = 'mean') → None

Initializes internal Module state, shared by both nn.Module and ScriptModule.

apply_aggregate(batch_size: int, offsets: Tensor, inverse_indices: Tensor, unqiue_embedding_vectors: Tensor) → Tensor

Aggregates the unique embedding vectors using the defined mode.

Parameters:

• batch_size (int) – Batch size of the input data.

• offsets (Tensor) – Offsets of inverse indices for each batch. Defines which embedding vectors are aggregated.

• inverse_indices (Tensor) – Flattened bag of indices for each batch.

• unqiue_embedding_vectors (Tensor) – Unique embedding vectors to be aggregated.

Returns:

The aggregated embedding vectors.

Return type:

Tensor

forward(indices: Tensor, offsets: Tensor) → Tensor

Generates embeddings from given tokens and offsets.

Parameters:

• indices (Tensor) – The tokens to be embedded.

• offsets (Tensor) – The offsets describing the starting points of batch items.

Returns:

The embedded and aggregated tokens.

Return type:

Tensor