quaterion_models.heads.sequential_head module

class SequentialHead(*args, output_size: int)

Bases: EncoderHead

A torch.nn.Sequential-like head layer that you can freely add any layers.

Unlike torch.nn.Sequential, it also expects the output size to be passed as a required keyword-only argument. It is required because some loss functions may need this information.

Parameters:

• args – Any sequence of torch.nn.Module instances. See torch.nn.Sequential for more info.

• output_size – Final output dimension from this head.

Examples:

head = SequentialHead(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 30),
    output_size=30
)

append(module: Module) → SequentialHead

forward(input_vectors: Tensor, meta=None) → Tensor

Forward pass for this head layer.

Just like torch.nn.Sequential, it passes the input to the first module, and the output of each module is input to the next. The final output of this head layer is the output from the last module in the sequence.

Parameters:

• input_vectors – Batch of input vectors.

• meta – Optional metadata for this batch.

Returns:

Output from the last module in the sequence.

get_config_dict() → Dict[str, Any]

Constructs savable params dict

Returns:

Serializable parameters for __init__ of the Module

classmethod load(input_path: str) → EncoderHead

save(output_path)

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