zchen/b2txt25
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Remove CPU optimization call and add logging for TPU strategy and data pipeline performance

Browse Source
This commit is contained in:
Zchen
2025-10-17 11:45:20 +08:00
parent eb4e3fc69f
commit 57360bec8a
1 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
model_training_nnn_tpu/trainer_tf.py
View File

@@ -50,9 +50,6 @@ class BrainToTextDecoderTrainerTF:
self.args = args self.args = args
self.logger = None self.logger = None
# Optimize CPU utilization for data pipeline (利用224核心)
self._configure_cpu_optimization()
# Initialize TPU strategy # Initialize TPU strategy
self.strategy = create_tpu_strategy() self.strategy = create_tpu_strategy()
if self.strategy is None: if self.strategy is None:
@@ -60,6 +57,8 @@ class BrainToTextDecoderTrainerTF:
print(f"Training on {self.strategy.num_replicas_in_sync} TPU cores") print(f"Training on {self.strategy.num_replicas_in_sync} TPU cores")
print(f"Strategy type: {type(self.strategy).__name__}") print(f"Strategy type: {type(self.strategy).__name__}")
print("💡 Using tf.data.AUTOTUNE for optimal data pipeline performance")
print("📝 Ensure create_input_fn uses AUTOTUNE for .map() and .prefetch() operations")
# Configure mixed precision for TPU v5e-8 # Configure mixed precision for TPU v5e-8
if args.get('use_amp', True): if args.get('use_amp', True):