File size: 1,660 Bytes
ab4d9a8 | 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 |
class MixerBlock(layers.Layer):
def __init__(self, seq_len, dim, token_mlp_dim, channel_mlp_dim, dropout=0.0):
super().__init__()
self.seq_len = seq_len
self.dim = dim
self.token_mlp_dim = token_mlp_dim
self.channel_mlp_dim = channel_mlp_dim
self.ln1 = layers.LayerNormalization(epsilon=1e-6, dtype=tf.float32)
# token-mixing MLP: operate over tokens => apply Dense on transposed axis
self.token_fc1 = layers.Dense(token_mlp_dim, activation='gelu', dtype=tf.float32)
self.token_fc2 = layers.Dense(seq_len, dtype=tf.float32)
self.ln2 = layers.LayerNormalization(epsilon=1e-6, dtype=tf.float32)
# channel-mixing MLP: operate per-token over channels
self.channel_fc1 = layers.Dense(channel_mlp_dim, activation='gelu', dtype=tf.float32)
self.channel_fc2 = layers.Dense(dim, dtype=tf.float32)
self.dropout = layers.Dropout(dropout)
def call(self, x, training=None):
# x: (B, L, D)
B = tf.shape(x)[0]
L = tf.shape(x)[1]
D = tf.shape(x)[2]
# Token-mixing
y = self.ln1(x) # (B, L, D)
y_t = tf.transpose(y, perm=[0,2,1]) # (B, D, L)
y_t = self.token_fc1(y_t) # (B, D, token_mlp_dim)
y_t = self.token_fc2(y_t) # (B, D, L)
y = tf.transpose(y_t, perm=[0,2,1]) # (B, L, D)
x = x + self.dropout(y, training=training)
# Channel-mixing
z = self.ln2(x)
z = self.channel_fc1(z)
z = self.channel_fc2(z)
x = x + self.dropout(z, training=training)
return x |