peal.architectures.basic_modules.ImgSelfAttentionLayer

class peal.architectures.basic_modules.ImgSelfAttentionLayer

Bases: Module

Self-attention over the spatial positions of an image feature map.

Two extra channels holding normalised row / column coordinates are concatenated to the input as positional encoding, the map is flattened to a (B, H*W, C+2) sequence, attended, and reshaped back to (B, C, H, W).

Parameters:

• inplanes (int) – Number of input channels C; the attention embedding is C + 2.

• num_heads (int, default 1) – Number of attention heads.

• use_masking (bool, default False) – Apply a causal mask over the flattened positions.

Methods

forward

Attend over spatial positions of x of shape (B, C, H, W).

__init__(inplanes, num_heads=1, use_masking=False)

Parameters:

• inplanes (int)

• num_heads (int)

• use_masking (bool)

Return type:

None

forward(x)

Attend over spatial positions of x of shape (B, C, H, W).

Returns:

Same shape as x; the coordinate channels are dropped again.

Return type:

torch.Tensor