File size: 2,348 Bytes
0f46930 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
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)
|