quaterion_models.heads.stacked_projection_head module

class StackedProjectionHead(input_embedding_size: int, output_sizes: List[int], activation_fn: str = 'relu', dropout: float = 0.0)

Bases: EncoderHead

Stacks any number of projection layers with specified output sizes.

Parameters:

• input_embedding_size – Dimensionality of the input to this stack of layers.

• output_sizes – List of output sizes for each one of the layers stacked.

• activation_fn – Name of the activation function to apply between the layers stacked. Must be an attribute of torch.nn.functional and defaults to relu.

• dropout – Probability of Dropout. If dropout > 0., apply dropout layer on embeddings before applying head layer transformations

get_config_dict() → Dict[str, Any]

Constructs savable params dict

Returns:

Serializable parameters for __init__ of the Module

transform(input_vectors: Tensor) → Tensor

Apply head-specific transformations to the embeddings tensor. Called as part of forward function, but with generic wrappings

Parameters:

input_vectors – Concatenated embeddings of all encoders. Shape: (batch_size, self.input_embedding_size)

Returns:

Final embeddings for a batch – (batch_size, self.output_size)

property output_size: int

training: bool