bitorch.layers.bembedding.BEmbedding

class bitorch.layers.bembedding.BEmbedding(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)

Binarized version of pytorchs embedding layer. 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_padding	Applies padding to the embedding vectors.
forward	Generates embeddings for received tokens.
select_unique_vectors	Given a flat tensor of indices, return the unique indices, their inverse in the original tensor, and a tensor with embedding vectors that are indexed by the unique indices.
set_optimizable_weights	Inject the weights to be optimized into the optimizer.
set_optimizer	Set the optimizer to set parameters to be optimized dynamically during training.
set_weight
step	Step the BEmbedding by copying the optimized unique embedding vectors into the binary embedding table.
transform_zeros	If the sign_bool property is set, replaces 0 with -1.
unique_wrapper	Compute the unique values and inverse indices of a given tensor.

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) → None

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

apply_padding(indices: Tensor, embedding_vectors: Tensor) → Tensor

Applies padding to the embedding vectors. Sets the embedding vector to zero where the given unique index matches the padding_idx property. This operation is inplace.

Parameters:

• indices (Tensor) – Indices of the embedding vectors.

• embedding_vectors (Tensor) – Embedding vectors to be padded.

Returns:

Padded embedding vectors.

Return type:

Tensor

forward(input: Tensor) → Tensor

Generates embeddings for received tokens.

Parameters:

input (Tensor) – indices for embedding

Returns:

embeddings for given token

Return type:

Tensor

select_unique_vectors(flat_indices: Tensor) → Tuple[Tensor, Tensor, Tensor]

Given a flat tensor of indices, return the unique indices, their inverse in the original tensor, and a tensor with embedding vectors that are indexed by the unique indices.

Parameters:

flat_indices (Tensor) – A flat tensor of indices that query the embedding table.

Returns:

unqiue indices, inverse indices, unique indexed embedding vectors

Return type:

Tuple[Tensor, Tensor, Tensor]

set_optimizable_weights(weights: Tensor) → None

Inject the weights to be optimized into the optimizer.

Parameters:

weights (Tensor) – The weights to be ioptimized.

set_optimizer(optimizer: Optimizer) → None

Set the optimizer to set parameters to be optimized dynamically during training.

Parameters:

optimizer (torch.optim.Optimizer) – The optimizer of the BEmbedding.

step() → None

Step the BEmbedding by copying the optimized unique embedding vectors into the binary embedding table.

transform_zeros(embedding_vectors: Tensor) → Tensor

If the sign_bool property is set, replaces 0 with -1. This operation is inplace.

Parameters:

embedding_vectors (Tensor) – The tensor to be modified.

Returns:

The modified input tensor

Return type:

Tensor

unique_wrapper(tensor: Tensor) → Tuple[Tensor, Tensor]

Compute the unique values and inverse indices of a given tensor. Uses numpy when on cpu and otherwise pytorch.

Parameters:

tensor (Tensor) – Tensor to compute the unique values from.

Returns:

unique values, inverse indices

Return type:

Tuple[Tensor, Tensor]