Update Mo.py

Mo.py

@@ -10,7 +10,8 @@ import tensorflow.keras.backend as K
 from tensorflow.keras import mixed_precision
 policy = mixed_precision.Policy('mixed_float16')  # fp16
 mixed_precision.set_global_policy(policy)
-print("✅ Mixed precision 적용:", policy)
+
+
 print('1')
 tf.get_logger().setLevel("ERROR")
 SEED = 42
@@ -18,18 +19,23 @@ tf.random.set_seed(SEED)
 np.random.seed(SEED)
 
 # TPU 초기화
-gpus = tf.config.list_physical_devices('GPU')
-if gpus:
-    try:
-        for gpu in gpus:
-            tf.config.experimental.set_memory_growth(gpu, True)
-        strategy = tf.distribute.MirroredStrategy(devices=[f"/GPU:{i}" for i in range(len(gpus))])
-        print(f"✅ GPU {len(gpus)}개 사용: {strategy.num_replicas_in_sync} replicas")
-    except RuntimeError as e:
-        print("⚠️ GPU 설정 에러:", e)
-else:
+try:
+    resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu="local")
+    tf.tpu.experimental.initialize_tpu_system(resolver)
+    strategy = tf.distribute.TPUStrategy(resolver)
+    print("✅ TPU 초기화 완료:", resolver.cluster_spec().as_dict())
+    on_tpu = True
+
+except Exception as e:
+    print("⚠️ TPU 미사용, GPU/CPU로 진행:", e)
     strategy = tf.distribute.get_strategy()
-    print("⚠️ GPU 없음, CPU 사용")
+    on_tpu = False
+
+# Mixed precision
+from tensorflow.keras import mixed_precision
+policy = mixed_precision.Policy("mixed_bfloat16" if on_tpu else "float32")
+mixed_precision.set_global_policy(policy)
+print("✅ Mixed precision:", policy)
 
 # =======================
 # 1) 파일 다운로드