Kaggle TPU 3.8 VM上TensorFlow推理返回批量大小，如何获取真实预测值？

问题原因 #

在TPUStrategy分布式环境中，直接调用model.predict()会默认对各TPU核心的输出执行归约聚合（如求和），导致返回的不是单样本的真实预测结果，而是聚合后的数值。

解决方法 #

通过strategy.run执行推理，并收集各TPU核心的输出结果进行拼接，即可获取真实的单样本预测值。具体步骤如下：

1. 确保模型在策略作用域内构建/加载
   你的初始化代码已正确创建TPU策略，需保证模型是在strategy.scope()中构建或加载的：

   import tensorflow as tf
   import numpy as np
   
   # 你的TPU初始化代码
   resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
   tf.config.experimental_connect_to_cluster(resolver)
   tf.tpu.experimental.initialize_tpu_system(resolver)
   strategy = tf.distribute.TPUStrategy(resolver)
   
   # 在策略作用域内构建/加载模型
   with strategy.scope():
       # 替换为你的模型构建或权重加载代码
       model = tf.keras.models.Sequential([...])
       # 如果是训练好的模型，加载权重
       # model.load_weights("your_model_weights.h5")
   

2. 分布式推理并收集结果
   定义推理函数，通过strategy.run执行，再拼接各核心的输出：

   # 定义推理步骤，关闭训练模式
   def predict_step(inputs):
       return model(inputs, training=False)
   
   # 将测试数据集转换为分布式数据集
   distributed_test_dataset = strategy.experimental_distribute_dataset(test_dataset)
   
   # 存储所有预测结果
   all_predictions = []
   
   # 取2个批次进行推理（对应你之前的take(2)）
   for batch in distributed_test_dataset.take(2):
       # 在各TPU核心上执行推理
       per_replica_preds = strategy.run(predict_step, args=(batch,))
       # 收集各核心的结果并拼接成完整批次
       batch_preds = tf.concat(tf.distribute.experimental_local_results(per_replica_preds), axis=0)
       all_predictions.append(batch_preds.numpy())
   
   # 合并所有批次的结果
   all_predictions = np.concatenate(all_predictions, axis=0)
   

   此时all_predictions即为每个样本的真实预测结果，形状与你的模型输出一致。

TPU推理核心代码要点 #

• 必须在TPUStrategy作用域内创建模型，确保模型参数分布到各TPU核心
• 使用strategy.experimental_distribute_dataset处理测试数据集，保证数据均匀分配到各核心
• 通过strategy.run执行推理逻辑，tf.distribute.experimental_local_results获取各核心的原始输出
• 拼接各核心结果得到完整的单样本预测集合

内容的提问来源于stack exchange，提问作者Alekk