Create HyperConv1D.py

HyperConv1D.py

@@ -0,0 +1,64 @@
+class HyperConv1D(layers.Layer):
+    def __init__(self, d_model, k=7, mem_size=64, hyper_dim=128, dropout=0.0):
+        super().__init__()
+        assert k % 2 == 1
+        self.k = k
+        self.d_model = d_model
+        self.mem_size = mem_size
+
+        # Input projection
+        self.input_proj = layers.Dense(d_model, name="input_proj")
+
+        # Local depthwise conv
+        self.local_conv = layers.DepthwiseConv1D(kernel_size=k, padding='same', activation='silu')
+        self.local_proj = layers.Dense(d_model, name="local_proj")
+
+        # Hypernetwork: global -> scale vector
+        self.hyper = tf.keras.Sequential([
+            layers.Dense(hyper_dim, activation='gelu'),
+            layers.Dense(d_model)
+        ], name="hyper")
+
+        # Associative memory
+        self.mem_keys = self.add_weight((mem_size, d_model), initializer='glorot_uniform', trainable=True)
+        self.mem_vals = self.add_weight((mem_size, d_model), initializer='glorot_uniform', trainable=True)
+        self.mem_proj = layers.Dense(d_model)
+
+        self.norm = layers.LayerNormalization()
+        self.attn_pool = layers.Dense(1)
+
+    def call(self, x):
+        x_in = x
+        x_dtype = x.dtype  # 입력 dtype 기억
+
+        # 1) input projection
+        x_proj = self.input_proj(x)
+        # memory와 연산 위해 dtype 통일
+        mem_dtype = self.mem_keys.dtype
+        x_proj = tf.cast(x_proj, mem_dtype)
+
+        # 2) local conv
+        out_local = self.local_conv(x_proj)
+        # hypernetwork scaling
+        global_z = self.attn_pool(x_proj)
+        global_z = tf.nn.softmax(global_z, axis=1)
+        global_z = tf.reduce_sum(x_proj * global_z, axis=1)
+
+        scale = tf.expand_dims(tf.nn.sigmoid(self.hyper(global_z)), 1)
+        out_local = out_local * scale
+        out_local = self.local_proj(out_local)
+
+
+        # 3) associative memory
+        sims = tf.matmul(x_proj, self.mem_keys, transpose_b=True) / tf.math.sqrt(tf.cast(self.d_model, mem_dtype))
+        attn = tf.nn.softmax(sims, axis=-1)
+        mem_read = tf.matmul(attn, self.mem_vals)
+        mem_read = self.mem_proj(mem_read)
+
+        # 4) fuse & residual
+        out = out_local + mem_read
+        out = self.norm(x_proj + out)
+        out = tf.nn.silu(out)
+
+        # 최종 출력 dtype 원래 입력 dtype으로 캐스트
+        return tf.cast(out, x_dtype)