attention_module_(sca_hmda).py

# -*- coding: utf-8 -*-
"""Attention Module (SCA-HMDA).ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1TtaTJPFNbxj-fuwCpQrqlDh2rdJgKahf
"""

# attention.py

import tensorflow as tf
from tensorflow.keras import layers

class SCA_HMDA(layers.Layer):
    def __init__(self, num_heads, embed_dim, sparsity_factor=0.5, **kwargs):
        super(SCA_HMDA, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.embed_dim = embed_dim
        self.sparsity_factor = sparsity_factor

        self.spatial_attention = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.temporal_attention = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)
        self.cross_attention = layers.MultiHeadAttention(num_heads=num_heads, key_dim=embed_dim)

        self.gate = layers.Dense(embed_dim, activation="sigmoid")
        self.ln1 = layers.LayerNormalization()
        self.ln2 = layers.LayerNormalization()
        self.ffn = tf.keras.Sequential([
            layers.Dense(embed_dim * 4, activation="gelu"),
            layers.Dense(embed_dim)
        ])

        self.local_context = layers.Conv1D(embed_dim, kernel_size=3, padding="same", activation="relu")
        self.global_context = layers.GlobalAveragePooling1D()

    def sparse_selection(self, inputs):
        attention_scores = tf.reduce_mean(inputs, axis=-1)
        num_elements = tf.cast(tf.shape(inputs)[1], tf.float32)
        k = tf.cast(tf.math.maximum(1.0, self.sparsity_factor * num_elements), tf.int32)

        top_k_values, top_k_indices = tf.math.top_k(attention_scores, k=k, sorted=False)
        mask = tf.scatter_nd(tf.expand_dims(top_k_indices, axis=-1), tf.ones_like(top_k_values, dtype=tf.float32), [tf.shape(inputs)[1]])
        weighted_inputs = inputs * tf.expand_dims(mask, axis=-1)
        return weighted_inputs

    def call(self, inputs, temporal_inputs=None, mask=None):
        sparse_inputs = self.sparse_selection(inputs)
        cross_attn_output = self.cross_attention(sparse_inputs, inputs, inputs, attention_mask=mask)

        spatial_attn_output = self.spatial_attention(sparse_inputs, sparse_inputs, sparse_inputs, attention_mask=mask)
        temporal_attn_output = self.temporal_attention(temporal_inputs, temporal_inputs, temporal_inputs, attention_mask=mask) if temporal_inputs is not None else spatial_attn_output

        gated_output = self.gate(spatial_attn_output + temporal_attn_output)
        local_features = self.local_context(gated_output)
        global_features = tf.expand_dims(self.global_context(gated_output), axis=1)
        fused_features = local_features + global_features

        return self.ln2(self.ln1(fused_features + cross_attn_output) + self.ffn(fused_features))