peal.architectures.basic_modules.DimensionSwitchAttentionLayer

class peal.architectures.basic_modules.DimensionSwitchAttentionLayer

Bases: Module

Cross-attention that maps a spatial feature map to output_size slots.

A fixed random lookup table of output_size query vectors attends over the flattened, coordinate-augmented input positions, so the spatial axis is replaced by a learned-query axis of fixed length. The table is created as a torch.autograd.Variable rather than an nn.Parameter, so it is neither trained nor stored in state_dict.

Parameters:

• output_size (int) – Number of query slots (length of the output sequence).

• num_hidden (int) – Embedding dimension of the attention layer.

• num_positional_encodings (int) – Number of coordinate channels the input will carry; added to the width of the lookup table.

Methods

forward

Cross-attend the lookup queries over the positions of x.

__init__(output_size, num_hidden, num_positional_encodings)

forward(x)

Cross-attend the lookup queries over the positions of x.

Parameters:

x (torch.Tensor) – Shape (B, C, *spatial).

Returns:

Shape (B, num_hidden - len(spatial), output_size) after the coordinate channels are removed and the axes transposed.

Return type:

torch.Tensor