Update V2.py

V2.py

@@ -128,20 +128,18 @@ ds = ds.batch(BATCH_SIZE, drop_remainder=True)
 ds = ds.map(lambda v1, v2: ((v1, v2), tf.zeros([BATCH_SIZE], dtype=tf.float32)), num_parallel_calls=tf.data.AUTOTUNE)
 ds = ds.prefetch(tf.data.AUTOTUNE)
 
+
+
 class MixerBlock(layers.Layer):
-    """
-    TPU / mixed-precision 친화적 MLP-Mixer 블록 (토큰-믹싱 + 채널-믹싱).
-    내부 연산은 float32로 수행하여 안정성 확보, 출력은 입력 dtype으로 복원.
-    """
     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
 
-        # LayerNorm은 float32로 안정화
         self.ln1 = layers.LayerNormalization(epsilon=1e-6, dtype=tf.float32)
         # token-mixing MLP: operate over tokens => apply Dense on transposed axis
-        # Dense도 float32로 강제
         self.token_fc1 = layers.Dense(token_mlp_dim, activation='gelu', dtype=tf.float32)
         self.token_fc2 = layers.Dense(seq_len, dtype=tf.float32)
 
@@ -153,31 +151,26 @@ class MixerBlock(layers.Layer):
         self.dropout = layers.Dropout(dropout)
 
     def call(self, x, training=None):
-        """
-        x: (B, L, D) — dtype can be bfloat16/float32 depending on policy.
-        내부 계산은 float32로 수행하고, 반환은 원래 x.dtype으로 캐스팅.
-        """
-        orig_dtype = x.dtype
-        # 안정적 연산을 위해 float32로 변환
-        x_f = tf.cast(x, tf.float32)  # (B, L, D)
-
-        # ---- Token-mixing (Dense on token axis) ----
-        y = self.ln1(x_f)                      # (B, L, D) in float32
-        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_f = x_f + self.dropout(y, training=training)
-
-        # ---- Channel-mixing (per-token MLP) ----
-        z = self.ln2(x_f)                      # (B, L, D)
-        z = self.channel_fc1(z)                # (B, L, channel_mlp_dim)
-        z = self.channel_fc2(z)                # (B, L, D)
-        x_f = x_f + self.dropout(z, training=training)
-
-        # 최종: 원래 dtype으로 복원 (mixed-precision 이득 유지)
-        return tf.cast(x_f, orig_dtype)
-
+        # 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
 
 class L2NormLayer(layers.Layer):
     def __init__(self, axis=1, epsilon=1e-10, **kwargs):
@@ -295,4 +288,4 @@ history = model.fit(ds, epochs=EPOCHS, steps_per_epoch=steps_per_epoch, verbose=
 
 # encoder 가중치 저장
 encoder.save_weights("encoder_fit.weights.h5")
-print("Training finished and weights saved.")
+print("Training finished and weights saved.")