TFAttentionBlocks Module

Provides helper functions and classes for building and utilizing attention blocks in TensorFlow models.

class HMB.TFAttentionBlocks.CoordAttention(*args, **kwargs)

Bases: Layer

Coordinate Attention block implementation as a Keras Layer.

This block factorizes channel attention into two directional attention maps that encode positional information along height and width respectively.

Parameters:

reduction (int) – Reduction ratio for the internal bottleneck. Default is 32.

Initialize CoordAttention.

Parameters:

reduction (int) – Reduction ratio for the bottleneck. Default is 32.

__init__(reduction=32, **kwargs)

Initialize CoordAttention.

Parameters:

reduction (int) – Reduction ratio for the bottleneck. Default is 32.

build(inputShape)

Build 1x1 conv layers used by the block.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass for coordinate attention.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after applying coordinate attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.SCSEBlock(*args, **kwargs)

Bases: Layer

Spatial-Channel Squeeze-and-Excitation (scSE) block implementation as a Keras Layer.

This block combines channel Squeeze-and-Excitation (cSE) and a spatial squeeze-and-excitation (sSE) to recalibrate features both channel-wise and spatially.

Parameters:

ratio (int) – Reduction ratio for the channel bottleneck. Default is 16.

Initialize the scSEBlock.

Parameters:

ratio (int) – Reduction ratio for the channel bottleneck. Default is 16.

__init__(ratio=16, **kwargs)

Initialize the scSEBlock.

Parameters:

ratio (int) – Reduction ratio for the channel bottleneck. Default is 16.

build(inputShape)

Build dense layers for cSE and a 1x1 conv for sSE.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: compute cSE and sSE branches and fuse their outputs.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after applying scSE attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.CBAMBlock(*args, **kwargs)

Bases: Layer

Convolutional Block Attention Module (CBAM) implementation as a Keras Layer. This layer applies channel and spatial attention mechanisms to refine feature maps.

Parameters:

• ratio (int) – Reduction ratio for channel attention MLP. Default is 8.

• kernelSize (int) – Kernel size for spatial attention convolution. Default is 7.

Examples

from HMB.TFAttentionBlocks import CBAMBlock
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, Input

model = Sequential([
  Input(shape=(64, 64, 128)),
  Conv2D(128, (3, 3), padding="same", activation="relu"),
  CBAMBlock(ratio=8, kernelSize=7),
  Conv2D(256, (3, 3), padding="same", activation="relu")
])
model.summary()

Initialize CBAM block with given parameters.

Parameters:

• ratio (int) – Reduction ratio for channel attention MLP. Default is 8.

• kernelSize (int) – Kernel size for spatial attention convolution. Default is 7.

__init__(ratio=8, kernelSize=7, **kwargs)

Initialize CBAM block with given parameters.

Parameters:

• ratio (int) – Reduction ratio for channel attention MLP. Default is 8.

• kernelSize (int) – Kernel size for spatial attention convolution. Default is 7.

build(inputShape)

Build the CBAM block components based on input shape.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: apply channel attention followed by spatial attention.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Refined feature map after applying CBAM attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.SEBlock(*args, **kwargs)

Bases: Layer

Squeeze-and-Excitation (SE) block implementation as a Keras Layer. This layer adaptively recalibrates channel-wise feature responses by explicitly modeling interdependencies between channels.

Parameters:

ratio (int) – Reduction ratio for the bottleneck in the SE block. Default is 16.

Examples

from HMB.TFAttentionBlocks import SEBlock
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, Input

model = Sequential([
  Input(shape=(64, 64, 128)),
  Conv2D(128, (3, 3), padding="same", activation="relu"),
  SEBlock(ratio=16),
  Conv2D(256, (3, 3), padding="same", activation="relu")
])
model.summary()

Initialize SE block with given parameters.

Parameters:

ratio (int) – Reduction ratio for the bottleneck in the SE block. Default is 16.

__init__(ratio=16, **kwargs)

Initialize SE block with given parameters.

Parameters:

ratio (int) – Reduction ratio for the bottleneck in the SE block. Default is 16.

build(inputShape)

Build the SE block components based on input shape.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: apply squeeze (global pooling) and excitation (MLP) to compute channel-wise weights and scale the input features.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Recalibrated feature map after applying SE attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.ECABlock(*args, **kwargs)

Bases: Layer

Efficient Channel Attention (ECA) block implementation as a Keras Layer. This layer captures local cross-channel interactions without dimensionality reduction.

Parameters:

• gamma (int) – Parameter to compute kernel size. Default is 2.

• b (int) – Parameter to compute kernel size. Default is 1.

Initialize ECA block with given parameters.

Parameters:

• gamma (int) – Parameter to compute kernel size. Default is 2.

• b (int) – Parameter to compute kernel size. Default is 1.

__init__(gamma=2, b=1, **kwargs)

Initialize ECA block with given parameters.

Parameters:

• gamma (int) – Parameter to compute kernel size. Default is 2.

• b (int) – Parameter to compute kernel size. Default is 1.

build(inputShape)

Build the ECA block components based on input shape. This method initializes the convolutional layer for channel attention.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: apply global pooling, 1D convolution for channel attention, and scale the input features.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Recalibrated feature map after applying ECA attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.MultiHeadSelfAttention(*args, **kwargs)

Bases: Layer

Multi-head self-attention adapted for 2D feature maps.

This layer flattens spatial dimensions (H*W) and applies a Transformer-style multi-head self-attention on the flattened sequence, then projects the attended features back to the original spatial shape.

Parameters:

• numHeads (int) – Number of attention heads. Default is 8.

• keyDim (int) – Dimensionality of each head’s key/query vectors. Default is 64.

Initialize MultiHeadSelfAttention.

Parameters:

• numHeads (int) – Number of attention heads.

• keyDim (int) – Dimension of each head.

__init__(numHeads=8, keyDim=64, **kwargs)

Initialize MultiHeadSelfAttention.

Parameters:

• numHeads (int) – Number of attention heads.

• keyDim (int) – Dimension of each head.

build(inputShape)

Build an internal tf.keras.layers.MultiHeadAttention instance. Input is expected to be (batch, H, W, C).

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: reshape (b,H,W,C) -> (b,seq_len,C), call the underlying MultiHeadAttention with query=value=key and reshape the output back.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after self-attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.NonLocalBlock(*args, **kwargs)

Bases: Layer

Non-local block (self-attention over feature map positions).

Computes pairwise affinities across spatial positions and aggregates global context. Useful to capture long-range dependencies in images.

Parameters:

interChannels (int|None) – Number of channels for the internal projections; if None it defaults to C//2 where C is input channels.

Initialize NonLocalBlock.

Parameters:

interChannels (int|None) – Internal projection channels.

__init__(interChannels=None, **kwargs)

Initialize NonLocalBlock.

Parameters:

interChannels (int|None) – Internal projection channels.

build(inputShape)

Build 1x1 conv projections for theta, phi and g, and an output 1x1 conv. inputShape is expected as (batch, H, W, C).

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: compute affinity matrix between positions and aggregate global context which is then projected and added to the input (residual).

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after non-local attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.BAMBlock(*args, **kwargs)

Bases: Layer

Bottleneck Attention Module (BAM) - lightweight channel + spatial attention.

This implementation builds a compact channel attention branch (MLP on pooled features) and a spatial branch (dilated conv stack) and multiplies them with the input feature map.

Parameters:

• reduction (int) – Channel reduction ratio for the channel branch. Default 16.

• dilationRates (tuple) – Dilation rates for the spatial conv stack. Default (1,2,4).

Initialize BAM block.

Parameters:

• reduction (int) – Reduction ratio for channel MLP.

• dilationRates (tuple) – Dilation rates for spatial convolutions.

__init__(reduction=16, dilationRates=(1, 2, 4), **kwargs)

Initialize BAM block.

Parameters:

• reduction (int) – Reduction ratio for channel MLP.

• dilationRates (tuple) – Dilation rates for spatial convolutions.

build(inputShape)

Build channel MLP and spatial dilated conv stack based on input channels.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: compute channel attention via pooled MLP and spatial attention via dilated convs, then apply both to input.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after BAM attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.GCBlock(*args, **kwargs)

Bases: Layer

Global Context (GC) block - simplified global attention for feature maps.

Computes a soft-attention mask over spatial positions, aggregates a global context vector and applies a small transform to produce a context tensor that is added back to the input (residual).

Parameters:

reduction (int) – Channel reduction for the transform MLP. Default 16.

Initialize GCBlock.

Parameters:

reduction (int) – Reduction factor for the internal transform.

__init__(reduction=16, **kwargs)

Initialize GCBlock.

Parameters:

reduction (int) – Reduction factor for the internal transform.

build(inputShape)

Build conv_mask and transform dense layers based on input channels.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: compute attention mask, aggregate weighted context and apply small transform then broadcast-add to the input.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after GC attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.AxialAttention(*args, **kwargs)

Bases: Layer

Axial attention: factorized attention applied along height then width axes. Applies a lightweight multi-head self-attention along one axis at a time to reduce computational cost compared to full 2D self-attention.

Parameters:

• numHeads (int) – Number of attention heads. Default 4.

• keyDim (int) – Per-head dimensionality. Default 32.

Initialize AxialAttention.

Parameters:

• numHeads (int) – Number of heads.

• keyDim (int) – Dimension per head.

__init__(numHeads=4, keyDim=32, **kwargs)

Initialize AxialAttention.

Parameters:

• numHeads (int) – Number of heads.

• keyDim (int) – Dimension per head.

build(inputShape)

Instantiate two MultiHeadSelfAttention modules (one for each axis).

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: apply attention along height and width axes and sum results.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after axial attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.AttentionAugmentedConv(*args, **kwargs)

Bases: Layer

Attention-augmented convolutional layer.

Combines a standard convolutional branch with a self-attention branch and concatenates their outputs along the channel axis.

Parameters:

• filters (int) – Total output filters (conv + attention channels combined).

• kernelSize (int) – Convolution kernel size for the conv branch.

• numHeads (int) – Number of attention heads for the attention branch.

• keyDim (int) – Per-head dimension for attention branch.

Initialize AttentionAugmentedConv.

Parameters:

• filters (int) – Total filters after concatenation.

• kernelSize (int) – Kernel size for conv branch.

• numHeads (int) – Number of attention heads.

• keyDim (int) – Per-head dimension.

__init__(filters, kernelSize=3, numHeads=4, keyDim=32, **kwargs)

Initialize AttentionAugmentedConv.

Parameters:

• filters (int) – Total filters after concatenation.

• kernelSize (int) – Kernel size for conv branch.

• numHeads (int) – Number of attention heads.

• keyDim (int) – Per-head dimension.

build(inputShape)

Build conv branch and attention branch. Note: user should ensure filters > numHeads * keyDim so the conv branch has a positive number of channels.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: run conv branch and attention branch then concatenate.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, filters) after concatenation of conv and attention branches.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.SKBlock(*args, **kwargs)

Bases: Layer

Selective Kernel (SK) block: adaptively fuse multiple convolutional branches.

Parameters:

• filters (int) – Number of output channels for each branch (usually equals input channels).

• M (int) – Number of branches (kernel choices). Default 2.

• G (int) – Grouping parameter (kept for API compatibility).

• r (int) – Reduction ratio for the selection MLP. Default 16.

Initialize SKBlock.

Parameters:

• filters (int) – Number of output channels per branch.

• M (int) – Number of branches.

• G (int) – Grouping parameter.

• r (int) – Reduction ratio for selection MLP.

__init__(filters, M=2, G=1, r=16, **kwargs)

Initialize SKBlock.

Parameters:

• filters (int) – Number of output channels per branch.

• M (int) – Number of branches.

• G (int) – Grouping parameter.

• r (int) – Reduction ratio for selection MLP.

build(inputShape)

Build convolutional branches and selection MLP.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: compute branch outputs, aggregate and apply attention-based fusion.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after selective kernel fusion.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.TripletAttention(*args, **kwargs)

Bases: Layer

Triplet attention: cross-dimension spatial attention using rotated inputs. Applies spatial attention on three views (original, HW-transposed, and original again — the current implementation applies the original view twice) and averages the three attended outputs.

Parameters:

kernelSize (int) – Kernel size for the spatial conv used to compute attention maps.

Initialize TripletAttention.

Parameters:

kernelSize (int) – Kernel size for attention conv.

__init__(kernelSize=7, **kwargs)

Initialize TripletAttention.

Parameters:

kernelSize (int) – Kernel size for attention conv.

build(inputShape)

Build the small conv used to compute spatial attention maps.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: apply spatial attention to original and transposed views, then average the results to obtain final output.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after triplet attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.PositionAttentionModule(*args, **kwargs)

Bases: Layer

Position Attention Module (PAM) used in DANet as a spatial attention block.

This module computes a spatial (position) affinity across all positions and aggregates context accordingly. It is lightweight and works on NHWC tensors.

Parameters:

interChannels (int|None) – Number of channels for the internal projections. If None it defaults to max(1, C//8) where C is input channels.

Initialize PositionAttentionModule.

Parameters:

interChannels (int|None) – Number of channels for internal projections.

__init__(interChannels=None, **kwargs)

Initialize PositionAttentionModule.

Parameters:

interChannels (int|None) – Number of channels for internal projections.

build(inputShape)

Build 1x1 conv projections for query, key and value, and an output conv.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass for PAM.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after position attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.ChannelAttentionModule(*args, **kwargs)

Bases: Layer

Channel Attention Module (CAM) used in DANet as a channel-wise attention block.

This module computes inter-channel affinities by treating spatial locations as the feature dimension and produces channel-wise context.

Initialize ChannelAttentionModule.

__init__(**kwargs)

Initialize ChannelAttentionModule.

build(inputShape)

Nothing heavy to build for CAM besides a learnable scale parameter.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass for CAM.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after channel attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.DANetBlock(*args, **kwargs)

Bases: Layer

Dual Attention Network (DANet) block combining position and channel attention.

This block applies two parallel branches: a Position Attention Module (PAM) and a Channel Attention Module (CAM). Their outputs are fused and returned.

Parameters:

interChannels (int|None) – Optional internal channel size for PAM.

Initialize DANetBlock.

Parameters:

interChannels (int|None) – Internal channel size for PAM.

__init__(interChannels=None, **kwargs)

Initialize DANetBlock.

Parameters:

interChannels (int|None) – Internal channel size for PAM.

build(inputShape)

Build internal PAM and CAM modules and small convs for feature transforms.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: apply branch transforms, attention modules and fuse outputs.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

class HMB.TFAttentionBlocks.CrissCrossAttention(*args, **kwargs)

Bases: Layer

Criss-Cross Attention (CCNet) layer implementation as a Keras Layer.

This module computes sparse dense contextual attention by attending along each position’s row and column separately, which reduces memory compared to full N x N attention while still capturing long-range dependencies.

Parameters:

interChannels (int|None) – Internal projection channels. If None it defaults to max(1, C//8) where C is the input channels.

Initialize CrissCrossAttention.

Parameters:

interChannels (int|None) – Number of channels for internal projections.

__init__(interChannels=None, **kwargs)

Initialize CrissCrossAttention.

Parameters:

interChannels (int|None) – Number of channels for internal projections.

build(inputShape)

Build 1x1 conv projections for query, key and value and an output conv.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass for criss-cross attention.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after criss-cross attention.

Return type:

tensorflow.Tensor

class HMB.TFAttentionBlocks.CrissCrossWrapperBlock(*args, **kwargs)

Bases: Layer

Wrapper block that applies CrissCrossAttention one or more times.

Parameters:

repeats (int) – Number of criss-cross attention passes to apply. Default is 1.

Initialize CrissCrossWrapperBlock.

Parameters:

• repeats (int) – Number of times to apply CrissCrossAttention.

• interChannels (int|None) – Internal projection channels passed to attention.

__init__(repeats=1, interChannels=None, **kwargs)

Initialize CrissCrossWrapperBlock.

Parameters:

• repeats (int) – Number of times to apply CrissCrossAttention.

• interChannels (int|None) – Internal projection channels passed to attention.

build(inputShape)

Instantiate the repeated CrissCrossAttention layers.

Parameters:

inputShape (tuple) – Shape of the input tensor.

call(inputs)

Forward pass: sequentially apply the CrissCrossAttention blocks.

Parameters:

inputs (tensorflow.Tensor) – Input feature map of shape (b, H, W, C).

Returns:

Output feature map of shape (b, H, W, C) after applying the repeated criss-cross attention blocks.

Return type:

tensorflow.Tensor