import math
import tensorflow as tf
from tensorflow.keras.layers import (
Layer, Dense, Conv2D, Conv1D, GlobalAveragePooling2D,
GlobalMaxPooling2D, Activation, Add, Multiply, Softmax,
Reshape, Permute, Concatenate, DepthwiseConv2D, MultiHeadAttention
)
from tensorflow.keras.utils import register_keras_serializable
[docs]
@register_keras_serializable(package="HMB")
class CBAMBlock(Layer):
r'''
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
--------
.. code-block:: python
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()
'''
# Store initialization args for later serialization.
[docs]
def __init__(self, ratio=8, kernelSize=7, **kwargs):
r'''
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.
'''
# Call parent initializer.
super(CBAMBlock, self).__init__(**kwargs)
# Set channel reduction ratio.
self.ratio = ratio
# Set kernel size for spatial attention.
self.kernelSize = kernelSize
# Build layer components when input shape is known.
[docs]
def build(self, inputShape):
r'''
Build the CBAM block components based on input shape.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
# inputShape is (batch, H, W, C).
# Determine number of channels.
channel = int(inputShape[-1])
# Compute bottleneck units.
hiddenUnits = max(1, channel // self.ratio)
self.sharedLayerOne = Dense(
hiddenUnits, # Reduced dimension.
activation="relu", # ReLU activation.
kernel_initializer="he_normal", # He normal initialization.
use_bias=True, # Use bias in dense layer.
name="cbam_shared_layer_one"
)
self.sharedLayerTwo = Dense(
channel, # Restore original channel dimension.
kernel_initializer="he_normal", # He normal initialization.
use_bias=True, # Use bias in dense layer.
name="cbam_shared_layer_two"
)
self.conv = Conv2D(
1, (self.kernelSize, self.kernelSize),
padding="same", # Ensure output size matches input size.
activation="sigmoid", # Sigmoid for attention map.
kernel_initializer="he_normal", # He normal initialization.
use_bias=False, # No bias needed.
)
# Create pooling layers.
self.gap = GlobalAveragePooling2D(name="cbam_global_avg_pool")
self.gmp = GlobalMaxPooling2D(name="cbam_global_max_pool")
# Finish build.
super(CBAMBlock, self).build(inputShape)
# Forward pass applying channel and spatial attention.
[docs]
def call(self, inputs):
r'''
Forward pass: apply channel attention followed by spatial attention.
Parameters:
inputs (tensorflow.Tensor): Input feature map of shape (b, H, W, C).
Returns:
tensorflow.Tensor: Refined feature map after applying CBAM attention.
'''
# Channel attention.
# Global average pooling across spatial dims.
avgPool = self.gap(inputs)
# Pass through shared MLP layer 1.
avgPool = self.sharedLayerOne(avgPool)
# Pass through shared MLP layer 2.
avgPool = self.sharedLayerTwo(avgPool)
# Global max pooling across spatial dims.
maxPool = self.gmp(inputs)
# Shared MLP layer 1 on max-pooled features.
maxPool = self.sharedLayerOne(maxPool)
# Shared MLP layer 2 on max-pooled features.
maxPool = self.sharedLayerTwo(maxPool)
# Sum channel attention contributions.
channelAtt = Add()([avgPool, maxPool])
# Sigmoid activation for channel weights.
channelAtt = Activation("sigmoid")(channelAtt)
# Reshape to (b,1,1,C).
channelAtt = tf.expand_dims(tf.expand_dims(channelAtt, axis=1), axis=1)
# Apply channel attention.
x = Multiply()([inputs, channelAtt])
# Spatial attention.
# Average across channels for spatial attention.
avgPoolSpatial = tf.reduce_mean(x, axis=-1, keepdims=True)
# Max across channels for spatial attention.
maxPoolSpatial = tf.reduce_max(x, axis=-1, keepdims=True)
# Concatenate spatial descriptors.
concat = tf.concat([avgPoolSpatial, maxPoolSpatial], axis=-1)
# Convolution to produce spatial attention map.
spatialAtt = self.conv(concat)
# Apply spatial attention.
out = Multiply()([x, spatialAtt])
return out # Return refined features.
[docs]
@register_keras_serializable(package="HMB")
class SEBlock(Layer):
r'''
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
--------
.. code-block:: python
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()
'''
# Store initialization args for later serialization.
[docs]
def __init__(self, ratio=16, **kwargs):
r'''
Initialize SE block with given parameters.
Parameters:
ratio (int): Reduction ratio for the bottleneck in the SE block. Default is 16.
'''
# Call parent initializer.
super(SEBlock, self).__init__(**kwargs)
# Set reduction ratio.
self.ratio = ratio
# Build layer components when input shape is known.
[docs]
def build(self, inputShape):
r'''
Build the SE block components based on input shape.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
# inputShape is (batch, H, W, C).
# Determine number of channels.
channel = int(inputShape[-1])
# Compute bottleneck units.
hiddenUnits = max(1, channel // self.ratio)
self.denseOne = Dense(
hiddenUnits, # Reduced dimension.
activation="relu", # ReLU activation.
kernel_initializer="he_normal", # He normal initialization.
use_bias=True, # Use bias in dense layer.
name="se_dense_one"
)
self.denseTwo = Dense(
channel, # Restore original channel dimension.
activation="sigmoid", # Sigmoid for attention weights.
kernel_initializer="he_normal", # He normal initialization.
use_bias=True, # Use bias in dense layer.
name="se_dense_two"
)
# Global average pooling layer.
self.gap = GlobalAveragePooling2D(name="se_global_avg_pool")
# Finish build.
super(SEBlock, self).build(inputShape)
# Forward pass applying squeeze and excitation.
[docs]
def call(self, inputs):
r'''
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:
tensorflow.Tensor: Recalibrated feature map after applying SE attention.
'''
# Squeeze: Global average pooling across spatial dims.
squeeze = self.gap(inputs)
# Excitation: First dense layer.
excitation = self.denseOne(squeeze)
# Excitation: Second dense layer.
excitation = self.denseTwo(excitation)
# Reshape to (b,1,1,C).
excitation = tf.expand_dims(tf.expand_dims(excitation, axis=1), axis=1)
# Scale input features by channel-wise weights.
out = Multiply()([inputs, excitation])
# Return recalibrated features.
return out
[docs]
@register_keras_serializable(package="HMB")
class ECABlock(Layer):
r'''
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.
'''
# Store initialization args for later serialization.
[docs]
def __init__(self, gamma=2, b=1, **kwargs):
r'''
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.
'''
# Call parent initializer.
super(ECABlock, self).__init__(**kwargs)
# Set gamma parameter.
self.gamma = gamma
# Set b parameter.
self.b = b
# Build layer components when input shape is known.
[docs]
def build(self, inputShape):
r'''
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.
'''
# inputShape is (batch, H, W, C).
# Determine number of channels.
channel = int(inputShape[-1])
# Compute kernel size using formula.
t = int(abs((math.log(channel, 2) + self.b) / self.gamma))
kSize = t if ((t % 2) == 1) else t + 1 # Ensure kernel size is odd.
# Use Conv1D acting over channel dimension: after global-pool we get
# (b, C) -> expanded to (b, C, 1) where C is the temporal dimension
# for Conv1D. This implements the ECA 1D conv over channels cleanly.
self.conv = Conv1D(
1, kernel_size=kSize,
padding="same",
activation="sigmoid",
kernel_initializer="he_normal",
use_bias=False,
name="eca_conv1d"
)
# Global average pooling layer.
self.gap = GlobalAveragePooling2D(name="eca_global_avg_pool")
super(ECABlock, self).build(inputShape) # Finish build.
# Forward pass applying efficient channel attention.
[docs]
def call(self, inputs):
r'''
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:
tensorflow.Tensor: Recalibrated feature map after applying ECA attention.
'''
# Global average pooling across spatial dims.
squeeze = self.gap(inputs)
# Reshape to (b, C, 1) so Conv1D can operate over the channel dimension
# (treated as time/steps). Conv1D returns (b, C, 1).
squeeze = tf.expand_dims(squeeze, axis=-1) # (b, C, 1).
excitation = self.conv(squeeze) # (b, C, 1)
# transpose to (b, 1, C) then expand to (b,1,1,C) to match inputs for
# broadcasting when scaling.
excitation = tf.transpose(excitation, perm=[0, 2, 1]) # (b,1,C)
excitation = tf.expand_dims(excitation, axis=1) # (b,1,1,C)
# Scale input features by channel-wise weights.
out = Multiply()([inputs, excitation])
# Return recalibrated features.
return out
[docs]
@register_keras_serializable(package="HMB")
class MultiHeadSelfAttention(Layer):
r'''
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.
'''
[docs]
def __init__(self, numHeads=8, keyDim=64, **kwargs):
r'''
Initialize MultiHeadSelfAttention.
Parameters:
numHeads (int): Number of attention heads.
keyDim (int): Dimension of each head.
'''
super(MultiHeadSelfAttention, self).__init__(**kwargs)
# Use camelCase for internal variable names as requested.
self.numHeads = numHeads
self.keyDim = keyDim
# create underlying keras MultiHeadAttention here so it's available
# before build/call (avoids NoneType callable issue in some call paths)
self.mha = MultiHeadAttention(num_heads=self.numHeads, key_dim=self.keyDim)
[docs]
def build(self, inputShape):
r'''
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.
'''
# Nothing else to build (self.mha created in __init__), keep for API.
super(MultiHeadSelfAttention, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
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:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after self-attention.
'''
b = tf.shape(inputs)[0]
h = tf.shape(inputs)[1]
w = tf.shape(inputs)[2]
c = tf.shape(inputs)[3]
seqLen = h * w
# Flatten spatial dims to sequence.
x = tf.reshape(inputs, (b, seqLen, c))
# Use the underlying multi-head attention; returns (b, seq_len, c)
out = self.mha(query=x, value=x, key=x)
out = tf.reshape(out, (b, h, w, c))
return out
[docs]
@register_keras_serializable(package="HMB")
class NonLocalBlock(Layer):
r'''
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.
'''
[docs]
def __init__(self, interChannels=None, **kwargs):
r'''
Initialize NonLocalBlock.
Parameters:
interChannels (int|None): Internal projection channels.
'''
super(NonLocalBlock, self).__init__(**kwargs)
self.interChannels = interChannels
[docs]
def build(self, inputShape):
r'''
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.
'''
_, h, w, c = inputShape
if self.interChannels is None:
self.interChannels = max(1, c // 2)
# Linear projections implemented as 1x1 convs.
self.theta = Conv2D(self.interChannels, 1, padding="same", use_bias=False, name="nlb_theta_conv")
self.phi = Conv2D(self.interChannels, 1, padding="same", use_bias=False, name="nlb_phi_conv")
self.g = Conv2D(self.interChannels, 1, padding="same", use_bias=False, name="nlb_g_conv")
self.outConv = Conv2D(c, 1, padding="same", use_bias=False, name="nlb_out_conv")
super(NonLocalBlock, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
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:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after non-local attention.
'''
b = tf.shape(inputs)[0]
h = tf.shape(inputs)[1]
w = tf.shape(inputs)[2]
c = tf.shape(inputs)[3]
thetaX = self.theta(inputs)
phiX = self.phi(inputs)
gX = self.g(inputs)
# flatten spatial dims: (b, N, ic).
thetaFlat = tf.reshape(thetaX, (b, -1, self.interChannels))
phiFlat = tf.reshape(phiX, (b, -1, self.interChannels))
gFlat = tf.reshape(gX, (b, -1, self.interChannels))
logits = tf.matmul(thetaFlat, phiFlat, transpose_b=True)
attn = tf.nn.softmax(logits, axis=-1)
out = tf.matmul(attn, gFlat)
out = tf.reshape(out, (b, h, w, self.interChannels))
out = self.outConv(out)
# residual connection.
return inputs + out
[docs]
@register_keras_serializable(package="HMB")
class BAMBlock(Layer):
r'''
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).
'''
[docs]
def __init__(self, reduction=16, dilationRates=(1, 2, 4), **kwargs):
r'''
Initialize BAM block.
Parameters:
reduction (int): Reduction ratio for channel MLP.
dilationRates (tuple): Dilation rates for spatial convolutions.
'''
super(BAMBlock, self).__init__(**kwargs)
self.reduction = reduction
self.dilationRates = dilationRates
[docs]
def build(self, inputShape):
r'''
Build channel MLP and spatial dilated conv stack based on input channels.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
c = int(inputShape[-1])
hidden = max(1, c // self.reduction)
# channel branch: small MLP.
self.channelFc1 = Dense(hidden, activation="relu", name="bam_channel_fc1")
self.channelFc2 = Dense(c, activation=None, name="bam_channel_fc2")
# spatial branch: stack of dilated convs.
self.spatialConvs = [
Conv2D(1, 3, padding="same", dilation_rate=d, activation=None, name=f"bam_spatial_conv_d{d}")
for d in self.dilationRates
]
self.sigmoid = Activation("sigmoid", name="bam_sigmoid")
self.globalAvgPool = GlobalAveragePooling2D(name="bam_global_avg_pool")
super(BAMBlock, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
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:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after BAM attention.
'''
# channel attention: global pooling -> MLP -> sigmoid -> reshape.
ch = self.globalAvgPool(inputs)
ch = self.channelFc1(ch)
ch = self.channelFc2(ch)
ch = self.sigmoid(ch)
ch = tf.expand_dims(tf.expand_dims(ch, axis=1), axis=1)
# spatial attention: pass input through dilated conv stack.
x = inputs
for conv in self.spatialConvs:
x = conv(x)
spat = self.sigmoid(x)
attn = Multiply()([inputs, ch])
attn = Multiply()([attn, spat])
return attn
[docs]
@register_keras_serializable(package="HMB")
class GCBlock(Layer):
r'''
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.
'''
[docs]
def __init__(self, reduction=16, **kwargs):
r'''
Initialize GCBlock.
Parameters:
reduction (int): Reduction factor for the internal transform.
'''
super(GCBlock, self).__init__(**kwargs)
self.reduction = reduction
[docs]
def build(self, inputShape):
r'''
Build conv_mask and transform dense layers based on input channels.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
c = int(inputShape[-1])
self.convMask = Conv2D(1, 1, padding="same", use_bias=False, name="gc_conv_mask")
self.softmax = Softmax(axis=1) # later applied over spatial positions.
self.transform = Dense(max(1, c // self.reduction), activation="relu", name="gc_transform_fc1")
self.transformOut = Dense(c, activation=None, name="gc_transform_fc2")
# Finish build.
super(GCBlock, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
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:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after GC attention.
'''
b = tf.shape(inputs)[0]
h = tf.shape(inputs)[1]
w = tf.shape(inputs)[2]
c = tf.shape(inputs)[3]
inputFlat = tf.reshape(inputs, (b, -1, c)) # (b, N, c).
mask = self.convMask(inputs) # (b,h,w,1).
maskFlat = tf.reshape(mask, (b, -1, 1)) # (b,N,1).
maskFlat = self.softmax(maskFlat) # softmax over N.
# context: weighted sum over positions -> (b,1,c) -> squeeze to (b,c).
context = tf.matmul(maskFlat, inputFlat, transpose_a=True)
context = tf.squeeze(context, axis=1)
contextTrans = self.transform(context)
contextTrans = self.transformOut(contextTrans)
contextTrans = tf.expand_dims(tf.expand_dims(contextTrans, axis=1), axis=1)
return inputs + contextTrans
[docs]
@register_keras_serializable(package="HMB")
class AxialAttention(Layer):
r'''
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.
'''
[docs]
def __init__(self, numHeads=4, keyDim=32, **kwargs):
r'''
Initialize AxialAttention.
Parameters:
numHeads (int): Number of heads.
keyDim (int): Dimension per head.
'''
super(AxialAttention, self).__init__(**kwargs)
self.numHeads = numHeads
self.keyDim = keyDim
[docs]
def build(self, inputShape):
r'''
Instantiate two MultiHeadSelfAttention modules (one for each axis).
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
# use the correct camelCase keyword expected by MultiHeadSelfAttention
self.mhsaH = MultiHeadSelfAttention(numHeads=self.numHeads, keyDim=self.keyDim)
self.mhsaW = MultiHeadSelfAttention(numHeads=self.numHeads, keyDim=self.keyDim)
super(AxialAttention, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
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:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after axial attention.
'''
# height-wise attention: transpose to bring height into the "sequence" axis.
xH = tf.transpose(inputs, perm=[0, 2, 1, 3]) # (b, W, H, C).
xH = self.mhsaH.call(xH)
xH = tf.transpose(xH, perm=[0, 2, 1, 3])
# width-wise attention (operate on original layout).
xW = self.mhsaW.call(inputs)
return xH + xW
[docs]
@register_keras_serializable(package="HMB")
class AttentionAugmentedConv(Layer):
r'''
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.
'''
[docs]
def __init__(self, filters, kernelSize=3, numHeads=4, keyDim=32, **kwargs):
r'''
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.
'''
super(AttentionAugmentedConv, self).__init__(**kwargs)
self.filters = filters
self.kernelSize = kernelSize
self.numHeads = numHeads
self.keyDim = keyDim
[docs]
def build(self, inputShape):
r'''
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.
'''
# conv branch output channels is total filters minus attention channels.
self.conv = Conv2D(
self.filters - self.keyDim * self.numHeads, self.kernelSize,
padding="same",
activation="relu",
name="aac_conv_branch"
)
# instantiate with correct kwarg name
self.mhsa = MultiHeadSelfAttention(numHeads=self.numHeads, keyDim=self.keyDim)
self.concat = Concatenate(axis=-1)
super(AttentionAugmentedConv, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
Forward pass: run conv branch and attention branch then concatenate.
Parameters:
inputs (tensorflow.Tensor): Input feature map of shape (b, H, W, C).
Returns:
tensorflow.Tensor: Output feature map of shape (b, H, W, filters) after concatenation of conv and attention branches.
'''
convOut = self.conv(inputs)
attnOut = self.mhsa.call(inputs)
out = self.concat([convOut, attnOut])
return out
[docs]
@register_keras_serializable(package="HMB")
class SKBlock(Layer):
r'''
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.
'''
[docs]
def __init__(self, filters, M=2, G=1, r=16, **kwargs):
r'''
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.
'''
super(SKBlock, self).__init__(**kwargs)
self.filters = filters
self.M = M
self.G = G
self.r = r
[docs]
def build(self, inputShape):
r'''
Build convolutional branches and selection MLP.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
# Build convolutional branches.
self.branches = []
for m in range(self.M):
kernelSize = 3 + 2 * m
conv = Conv2D(self.filters, kernelSize, padding="same", activation="relu")
self.branches.append(conv)
# Selection MLP.
self.globalPool = GlobalAveragePooling2D()
hiddenUnits = max(1, self.filters // self.r)
self.fc1 = Dense(hiddenUnits, activation="relu", name="sk_fc1")
self.fcs = [
Dense(self.filters, activation=None, name=f"sk_fc_branch_{m}")
for m in range(self.M)
]
# Softmax across branches (axis=1 corresponds to M after stacking).
self.softmax = Softmax(axis=1)
super(SKBlock, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
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:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after selective kernel fusion.
'''
# Compute each branch output -> list of tensors (b,H,W,C).
branchOuts = [conv(inputs) for conv in self.branches]
# Stack branch outputs -> (b, H, W, C, M).
stacked = tf.stack(branchOuts, axis=-1)
# Aggregate across branches: U shape (b,H,W,C).
U = tf.reduce_sum(stacked, axis=-1)
# Channel descriptor: (b, C).
s = self.globalPool(U)
# Shared FC -> (b, hidden).
z = self.fc1(s)
# Per-branch logits -> list of (b, C).
logitsList = [fc(z) for fc in self.fcs]
# Stack logits -> (b, M, C).
logits = tf.stack(logitsList, axis=1)
# Softmax over branches -> (b, M, C).
attn = self.softmax(logits)
# Reorder to (b, 1, 1, C, M) to broadcast with stacked (b,H,W,C,M).
attn = tf.transpose(attn, perm=[0, 2, 1]) # (b, C, M).
attn = tf.expand_dims(tf.expand_dims(attn, axis=1), axis=1) # (b,1,1,C,M).
# Weighted sum over branches -> (b,H,W,C).
V = tf.reduce_sum(stacked * attn, axis=-1)
return V
[docs]
@register_keras_serializable(package="HMB")
class TripletAttention(Layer):
r'''
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.
'''
[docs]
def __init__(self, kernelSize=7, **kwargs):
r'''
Initialize TripletAttention.
Parameters:
kernelSize (int): Kernel size for attention conv.
'''
super(TripletAttention, self).__init__(**kwargs)
self.kernelSize = kernelSize
[docs]
def build(self, inputShape):
r'''
Build the small conv used to compute spatial attention maps.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
self.conv = Conv2D(1, self.kernelSize, padding="same", activation="sigmoid", name="triplet_attention_conv")
super(TripletAttention, self).build(inputShape)
def _spatial_att(self, x):
r'''
Compute a spatial attention map from per-pixel descriptors
(avg and max across channels) and apply it to the input tensor.
Parameters:
x (tensorflow.Tensor): Input feature map of shape (b, H, W, C).
Returns:
tensorflow.Tensor: Spatially attended feature map of shape (b, H, W, C).
'''
avg = tf.reduce_mean(x, axis=-1, keepdims=True)
mx = tf.reduce_max(x, axis=-1, keepdims=True)
cat = tf.concat([avg, mx], axis=-1)
att = self.conv(cat)
return x * att
[docs]
def call(self, inputs):
r'''
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:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after triplet attention.
'''
a = self._spatial_att(inputs)
# HW-transposed view
p = tf.transpose(inputs, perm=[0, 2, 1, 3])
b = self._spatial_att(p)
b = tf.transpose(b, perm=[0, 2, 1, 3])
# Channel-transposed view: bring channels into the H position and apply
# spatial attention over that view, then transpose back. This yields a
# distinct third branch compared to the original implementation.
q = tf.transpose(inputs, perm=[0, 3, 2, 1])
c = self._spatial_att(q)
c = tf.transpose(c, perm=[0, 3, 2, 1])
out = (a + b + c) / 3.0
return out
[docs]
@register_keras_serializable(package="HMB")
class CoordAttention(Layer):
r'''
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.
'''
[docs]
def __init__(self, reduction=32, **kwargs):
r'''
Initialize CoordAttention.
Parameters:
reduction (int): Reduction ratio for the bottleneck. Default is 32.
'''
super(CoordAttention, self).__init__(**kwargs)
self.reduction = reduction
[docs]
def build(self, inputShape):
r'''
Build 1x1 conv layers used by the block.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
# inputShape is (batch, H, W, C).
channel = int(inputShape[-1])
hidden = max(1, channel // self.reduction)
# Shared transform after concatenating directional descriptors.
self.conv1 = Conv2D(
hidden, 1, padding="same", activation="relu", use_bias=True, name="coord_conv1"
)
# Produce height and width attention maps that will be broadcast-multiplied.
self.convH = Conv2D(channel, 1, padding="same", activation="sigmoid", use_bias=True, name="coord_conv_h")
self.convW = Conv2D(channel, 1, padding="same", activation="sigmoid", use_bias=True, name="coord_conv_w")
super(CoordAttention, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
Forward pass for coordinate attention.
Parameters:
inputs (tensorflow.Tensor): Input feature map of shape (b, H, W, C).
Returns:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after applying coordinate attention.
'''
# Compute directional pooled descriptors.
# Average along width -> shape (b, H, 1, C).
xH = tf.reduce_mean(inputs, axis=2, keepdims=True)
# Average along height -> shape (b, 1, W, C).
xW = tf.reduce_mean(inputs, axis=1, keepdims=True)
# Transpose xW to shape (b, W, 1, C) so it can be concatenated with xH.
xW = tf.transpose(xW, perm=[0, 2, 1, 3])
# Concatenate along the spatial (height) axis -> (b, H+W, 1, C).
cat = tf.concat([xH, xW], axis=1)
# Shared 1x1 conv to mix channel information -> (b, H+W, 1, hidden).
x = self.conv1(cat)
# Split back into height and width descriptors using the sizes from xH and xW.
xHPooled, xWPooled = tf.split(x, [tf.shape(xH)[1], tf.shape(xW)[1]], axis=1)
# Produce attention maps and reshape width attention back to (b, 1, W, C).
aH = self.convH(xHPooled)
aW = self.convW(xWPooled)
aW = tf.transpose(aW, perm=[0, 2, 1, 3])
# Broadcast multiply inputs by the two attention maps.
out = inputs * aH * aW
return out
[docs]
@register_keras_serializable(package="HMB")
class SCSEBlock(Layer):
r'''
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.
'''
[docs]
def __init__(self, ratio=16, **kwargs):
r'''
Initialize the scSEBlock.
Parameters:
ratio (int): Reduction ratio for the channel bottleneck. Default is 16.
'''
super(SCSEBlock, self).__init__(**kwargs)
self.ratio = ratio
[docs]
def build(self, inputShape):
r'''
Build dense layers for cSE and a 1x1 conv for sSE.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
channel = int(inputShape[-1])
hiddenUnits = max(1, channel // self.ratio)
# Channel SE MLP.
self.cseDense1 = Dense(
hiddenUnits, activation="relu",
kernel_initializer="he_normal",
use_bias=True,
name="scse_cse_dense1"
)
self.cseDense2 = Dense(
channel, activation="sigmoid",
kernel_initializer="he_normal",
use_bias=True,
name="scse_cse_dense2"
)
self.globalAvgPool = GlobalAveragePooling2D(name="scse_global_avg_pool")
# Spatial SE conv.
self.sseConv = Conv2D(
1, 1, padding="same",
activation="sigmoid",
kernel_initializer="he_normal",
use_bias=True,
name="scse_sse_conv"
)
super(SCSEBlock, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
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:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after applying scSE attention.
'''
# cSE branch: global pooling -> MLP -> reshape and scale.
ch = self.globalAvgPool(inputs)
ch = self.cseDense1(ch)
ch = self.cseDense2(ch)
ch = tf.expand_dims(tf.expand_dims(ch, axis=1), axis=1)
cSEOut = Multiply()([inputs, ch])
# sSE branch: 1x1 conv -> scale.
s = self.sseConv(inputs)
sSEOut = Multiply()([inputs, s])
# Fuse channel and spatial recalibrations by elementwise addition.
out = Add()([cSEOut, sSEOut])
return out
[docs]
@register_keras_serializable(package="HMB")
class PositionAttentionModule(Layer):
r'''
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.
'''
[docs]
def __init__(self, interChannels=None, **kwargs):
r'''
Initialize PositionAttentionModule.
Parameters:
interChannels (int|None): Number of channels for internal projections.
'''
super(PositionAttentionModule, self).__init__(**kwargs)
self.interChannels = interChannels
[docs]
def build(self, inputShape):
r'''
Build 1x1 conv projections for query, key and value, and an output conv.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
channel = int(inputShape[-1])
if (self.interChannels is None):
self.interChannels = max(1, channel // 8)
# 1x1 convs for query, key and value projections.
self.queryConv = Conv2D(self.interChannels, 1, padding="same", use_bias=False, name="pam_query_conv")
self.keyConv = Conv2D(self.interChannels, 1, padding="same", use_bias=False, name="pam_key_conv")
self.valueConv = Conv2D(self.interChannels, 1, padding="same", use_bias=False, name="pam_value_conv")
# Output projection to restore channel dimension.
self.outConv = Conv2D(channel, 1, padding="same", use_bias=False, name="pam_out_conv")
# Learnable scale parameter initialized to zero so the block can start as identity.
self.gamma = self.add_weight(name="pam_gamma", shape=(), initializer="zeros", trainable=True)
super(PositionAttentionModule, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
Forward pass for PAM.
Parameters:
inputs (tensorflow.Tensor): Input feature map of shape (b, H, W, C).
Returns:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after position attention.
'''
b = tf.shape(inputs)[0]
h = tf.shape(inputs)[1]
w = tf.shape(inputs)[2]
c = tf.shape(inputs)[3]
n = h * w # number of positions.
# Project inputs to query, key and value spaces.
projQuery = self.queryConv(inputs) # (b, H, W, inter).
projQuery = tf.reshape(projQuery, (b, -1, self.interChannels)) # (b, N, inter).
projKey = self.keyConv(inputs)
projKey = tf.reshape(projKey, (b, -1, self.interChannels)) # (b, N, inter).
projKeyT = tf.transpose(projKey, perm=[0, 2, 1]) # (b, inter, N).
# Affinity across positions.
energy = tf.matmul(projQuery, projKeyT) # (b, N, N).
attention = tf.nn.softmax(energy, axis=-1) # (b, N, N).
# Aggregate values using attention.
projValue = self.valueConv(inputs)
projValue = tf.reshape(projValue, (b, -1, self.interChannels)) # (b, N, inter).
out = tf.matmul(attention, projValue) # (b, N, inter).
out = tf.reshape(out, (b, h, w, self.interChannels)) # (b, H, W, inter).
out = self.outConv(out) # restore to (b, H, W, C).
# Residual connection with a learnable scale.
return inputs + self.gamma * out
[docs]
@register_keras_serializable(package="HMB")
class ChannelAttentionModule(Layer):
r'''
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.
'''
[docs]
def __init__(self, **kwargs):
r'''
Initialize ChannelAttentionModule.
'''
super(ChannelAttentionModule, self).__init__(**kwargs)
[docs]
def build(self, inputShape):
r'''
Nothing heavy to build for CAM besides a learnable scale parameter.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
channel = int(inputShape[-1])
self.gamma = self.add_weight(name="cam_gamma", shape=(), initializer="zeros", trainable=True)
super(ChannelAttentionModule, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
Forward pass for CAM.
Parameters:
inputs (tensorflow.Tensor): Input feature map of shape (b, H, W, C).
Returns:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after channel attention.
'''
b = tf.shape(inputs)[0]
h = tf.shape(inputs)[1]
w = tf.shape(inputs)[2]
c = tf.shape(inputs)[3]
n = h * w # number of positions.
# Reshape to (b, N, C) where N = H*W.
projQuery = tf.reshape(inputs, (b, -1, c)) # (b, N, C).
projKey = tf.transpose(projQuery, perm=[0, 2, 1]) # (b, C, N).
# Compute channel affinity matrix.
energy = tf.matmul(projKey, projQuery) # (b, C, C).
# Normalize affinities to probabilities.
attention = tf.nn.softmax(energy, axis=-1) # (b, C, C).
# Apply attention to the value (also channels viewed across spatial positions).
projValue = tf.transpose(projQuery, perm=[0, 2, 1]) # (b, C, N).
out = tf.matmul(attention, projValue) # (b, C, N).
out = tf.transpose(out, perm=[0, 2, 1]) # (b, N, C).
out = tf.reshape(out, (b, h, w, c)) # (b, H, W, C).
# Residual connection with learnable scale.
return inputs + self.gamma * out
[docs]
@register_keras_serializable(package="HMB")
class DANetBlock(Layer):
r'''
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.
'''
[docs]
def __init__(self, interChannels=None, **kwargs):
r'''
Initialize DANetBlock.
Parameters:
interChannels (int|None): Internal channel size for PAM.
'''
super(DANetBlock, self).__init__(**kwargs)
self.interChannels = interChannels
[docs]
def build(self, inputShape):
r'''
Build internal PAM and CAM modules and small convs for feature transforms.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
channel = int(inputShape[-1])
# Transform convs used before feeding branches.
self.convPAM = Conv2D(channel, 1, padding="same", use_bias=False, name="danet_conv_pam")
self.convCAM = Conv2D(channel, 1, padding="same", use_bias=False, name="danet_conv_cam")
# Instantiate PAM and CAM with appropriate internal channel size.
self.pam = PositionAttentionModule(self.interChannels)
self.cam = ChannelAttentionModule()
# Fuse conv to merge branch outputs.
self.fuseConv = Conv2D(channel, 1, padding="same", use_bias=False, name="danet_fuse_conv")
super(DANetBlock, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
Forward pass: apply branch transforms, attention modules and fuse outputs.
Parameters:
inputs (tensorflow.Tensor): Input feature map of shape (b, H, W, C).
'''
# Prepare features for each branch.
featP = self.convPAM(inputs)
featC = self.convCAM(inputs)
# Apply attention modules.
outP = self.pam.call(featP)
outC = self.cam.call(featC)
# Fuse outputs by summation and a final conv.
out = Add()([outP, outC])
out = self.fuseConv(out)
return out
[docs]
@register_keras_serializable(package="HMB")
class CrissCrossAttention(Layer):
r'''
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.
'''
[docs]
def __init__(self, interChannels=None, **kwargs):
r'''
Initialize CrissCrossAttention.
Parameters:
interChannels (int|None): Number of channels for internal projections.
'''
super(CrissCrossAttention, self).__init__(**kwargs)
self.interChannels = interChannels
[docs]
def build(self, inputShape):
r'''
Build 1x1 conv projections for query, key and value and an output conv.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
channel = int(inputShape[-1])
if self.interChannels is None:
self.interChannels = max(1, channel // 8)
# 1x1 convs for query, key and value projections.
self.queryConv = Conv2D(self.interChannels, 1, padding="same", use_bias=False, name="cc_query_conv")
self.keyConv = Conv2D(self.interChannels, 1, padding="same", use_bias=False, name="cc_key_conv")
self.valueConv = Conv2D(self.interChannels, 1, padding="same", use_bias=False, name="cc_value_conv")
# Output projection to restore channel dimension.
self.outConv = Conv2D(channel, 1, padding="same", use_bias=False, name="cc_out_conv")
# Learnable scale initialized to zero so the layer can start as identity.
self.gamma = self.add_weight(name="cc_gamma", shape=(), initializer="zeros", trainable=True)
super(CrissCrossAttention, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
Forward pass for criss-cross attention.
Parameters:
inputs (tensorflow.Tensor): Input feature map of shape (b, H, W, C).
Returns:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after criss-cross attention.
'''
b = tf.shape(inputs)[0]
h = tf.shape(inputs)[1]
w = tf.shape(inputs)[2]
# Project to query, key and value spaces.
projQ = self.queryConv(inputs) # (b, H, W, ic).
projK = self.keyConv(inputs) # (b, H, W, ic).
projV = self.valueConv(inputs) # (b, H, W, ic).
# Row-wise attention: treat each row as a sequence of length W.
qRow = tf.reshape(projQ, (-1, w, self.interChannels)) # (b*H, W, ic).
kRow = tf.reshape(projK, (-1, w, self.interChannels)) # (b*H, W, ic).
vRow = tf.reshape(projV, (-1, w, self.interChannels)) # (b*H, W, ic).
energyRow = tf.matmul(qRow, kRow, transpose_b=True) # (b*H, W, W).
attnRow = tf.nn.softmax(energyRow, axis=-1) # (b*H, W, W).
outRow = tf.matmul(attnRow, vRow) # (b*H, W, ic).
outRow = tf.reshape(outRow, (b, h, w, self.interChannels)) # (b, H, W, ic).
# Column-wise attention: treat each column as a sequence of length H.
# Transpose to bring columns into the sequence axis.
qCol = tf.transpose(projQ, perm=[0, 2, 1, 3]) # (b, W, H, ic).
kCol = tf.transpose(projK, perm=[0, 2, 1, 3]) # (b, W, H, ic).
vCol = tf.transpose(projV, perm=[0, 2, 1, 3]) # (b, W, H, ic).
qCol = tf.reshape(qCol, (-1, h, self.interChannels)) # (b*W, H, ic).
kCol = tf.reshape(kCol, (-1, h, self.interChannels)) # (b*W, H, ic).
vCol = tf.reshape(vCol, (-1, h, self.interChannels)) # (b*W, H, ic).
energyCol = tf.matmul(qCol, kCol, transpose_b=True) # (b*W, H, H).
attnCol = tf.nn.softmax(energyCol, axis=-1) # (b*W, H, H).
outCol = tf.matmul(attnCol, vCol) # (b*W, H, ic).
outCol = tf.reshape(outCol, (b, w, h, self.interChannels)) # (b, W, H, ic).
outCol = tf.transpose(outCol, perm=[0, 2, 1, 3]) # (b, H, W, ic).
# Aggregate row and column context.
out = outRow + outCol # (b, H, W, ic).
# Project back to input channel dimension and apply residual with scale.
out = self.outConv(out)
return inputs + self.gamma * out
[docs]
@register_keras_serializable(package="HMB")
class CrissCrossWrapperBlock(Layer):
r'''
Wrapper block that applies CrissCrossAttention one or more times.
Parameters:
repeats (int): Number of criss-cross attention passes to apply. Default is 1.
'''
[docs]
def __init__(self, repeats=1, interChannels=None, **kwargs):
r'''
Initialize CrissCrossWrapperBlock.
Parameters:
repeats (int): Number of times to apply CrissCrossAttention.
interChannels (int|None): Internal projection channels passed to attention.
'''
super(CrissCrossWrapperBlock, self).__init__(**kwargs)
self.repeats = repeats
self.interChannels = interChannels
[docs]
def build(self, inputShape):
r'''
Instantiate the repeated CrissCrossAttention layers.
Parameters:
inputShape (tuple): Shape of the input tensor.
'''
self.blocks = [CrissCrossAttention(self.interChannels) for _ in range(self.repeats)]
super(CrissCrossWrapperBlock, self).build(inputShape)
[docs]
def call(self, inputs):
r'''
Forward pass: sequentially apply the CrissCrossAttention blocks.
Parameters:
inputs (tensorflow.Tensor): Input feature map of shape (b, H, W, C).
Returns:
tensorflow.Tensor: Output feature map of shape (b, H, W, C) after applying the repeated criss-cross attention blocks.
'''
x = inputs
for blk in self.blocks:
x = blk.call(x)
return x
if __name__ == "__main__":
# Simple test case to verify layer instantiation.
import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, Input
inputShape = (64, 64, 128)
model = Sequential([
Input(shape=inputShape),
Conv2D(128, (3, 3), padding="same", activation="relu"),
CBAMBlock(ratio=8, kernelSize=7),
SEBlock(ratio=16),
ECABlock(gamma=2, b=1),
MultiHeadSelfAttention(numHeads=8, keyDim=64),
NonLocalBlock(),
BAMBlock(reduction=16, dilationRates=(1, 2, 4)),
GCBlock(reduction=16),
AxialAttention(numHeads=4, keyDim=32),
AttentionAugmentedConv(filters=256, kernelSize=3, numHeads=4, keyDim=32),
SKBlock(filters=128, M=2, G=1, r=16),
SKBlock(filters=256, M=2, G=1, r=16), # <- changed from 128 to 256
TripletAttention(kernelSize=7),
CoordAttention(reduction=32),
SCSEBlock(ratio=16),
Conv2D(256, (3, 3), padding="same", activation="relu")
])
model.summary()
# Test with random input.
x = np.random.rand(1, 64, 64, 128).astype(np.float32)
y = model.predict(x)
print("Output shape:", y.shape)
# Store model to disk and load it back to verify serialization.
model.save("test_attention_blocks_model.keras")
from tensorflow.keras.models import load_model
from HMB.TFAttentionBlocks import (
CBAMBlock, SEBlock, ECABlock, MultiHeadSelfAttention, NonLocalBlock,
BAMBlock, GCBlock, AxialAttention, AttentionAugmentedConv, SKBlock, TripletAttention,
PositionAttentionModule, ChannelAttentionModule, DANetBlock, CrissCrossAttention, CrissCrossWrapperBlock
)
customObjects = {
"CBAMBlock" : CBAMBlock,
"SEBlock" : SEBlock,
"ECABlock" : ECABlock,
"MultiHeadSelfAttention" : MultiHeadSelfAttention,
"NonLocalBlock" : NonLocalBlock,
"BAMBlock" : BAMBlock,
"GCBlock" : GCBlock,
"AxialAttention" : AxialAttention,
"AttentionAugmentedConv" : AttentionAugmentedConv,
"SKBlock" : SKBlock,
"TripletAttention" : TripletAttention,
"CoordAttention" : CoordAttention,
"SCSEBlock" : SCSEBlock,
"PositionAttentionModule": PositionAttentionModule,
"ChannelAttentionModule" : ChannelAttentionModule,
"DANetBlock" : DANetBlock,
"CrissCrossAttention" : CrissCrossAttention,
"CrissCrossWrapperBlock" : CrissCrossWrapperBlock,
}
# Load the model back.
loadedModel = load_model(
"test_attention_blocks_model.keras",
compile=True,
custom_objects=customObjects
)
# Verify loaded model works.
y2 = loadedModel.predict(x)
print("Loaded model output shape:", y2.shape)