我来帮你搞定这个问题！把Keras模型转成tf.Estimator后没用到多GPU，核心是没正确配置分布式策略和Estimator的运行参数。下面咱们一步步来实现分布式训练：

一、先选对分布式策略 #

TensorFlow提供了几种适配不同场景的分布式训练策略，根据你的硬件环境选择：

• MirroredStrategy：适合单机多GPU场景，会在每个GPU上复制模型，同步进行梯度更新，是单机分布式最常用的方案。
• MultiWorkerMirroredStrategy：适合多机多GPU场景，跨机器同步模型参数和梯度。

二、配置Estimator的RunConfig #

这一步是关键，要把分布式策略传给Estimator的运行配置，让它明确启用分布式训练。

举个单机多GPU的代码示例：

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense

# 1. 定义你的Keras模型
def build_keras_model():
    model = Sequential([
        Dense(64, activation='relu', input_shape=(10,)),
        Dense(64, activation='relu'),
        Dense(1, activation='sigmoid')
    ])
    model.compile(optimizer='adam',
                  loss='binary_crossentropy',
                  metrics=['accuracy'])
    return model

# 2. 初始化分布式策略
strategy = tf.distribute.MirroredStrategy()
print(f"当前使用 {strategy.num_replicas_in_sync} 块GPU进行训练")

# 3. 配置Estimator的运行参数
run_config = tf.estimator.RunConfig(
    train_distribute=strategy,  # 指定训练用的分布式策略
    model_dir="./distributed_model",  # 模型保存的路径
    log_step_count_steps=100  # 每100步打印一次训练日志
)

# 4. 把Keras模型转成tf.Estimator
estimator = tf.keras.estimator.model_to_estimator(
    keras_model=build_keras_model(),
    config=run_config
)

三、编写适配分布式的输入函数 #

输入函数要注意批量大小的设置：总批量大小应该是单GPU的批量乘以GPU数量，这样每个GPU处理的批量和单GPU训练时一致，保证训练效果的一致性。

def train_input_fn():
    # 生成示例训练数据（实际替换成你的业务数据集）
    x = tf.random.normal(shape=(10000, 10))
    y = tf.random.uniform(shape=(10000, 1), minval=0, maxval=2, dtype=tf.int32)
    
    # 构建Dataset，注意batch_size要乘以GPU数量
    dataset = tf.data.Dataset.from_tensor_slices((x, y))
    dataset = dataset.shuffle(buffer_size=10000) \
                     .repeat() \
                     .batch(32 * strategy.num_replicas_in_sync)
    return dataset

四、启动分布式训练 #

直接调用Estimator的train方法即可，分布式策略会自动处理GPU间的参数同步：

estimator.train(input_fn=train_input_fn, steps=2000)

多机多GPU的额外配置 #

如果是多机场景，需要在每台机器上设置TF_CONFIG环境变量，告诉TensorFlow集群的结构：
比如主节点（index=0）的环境变量：

export TF_CONFIG='{"cluster": {"worker": ["主节点IP:端口", "从节点IP:端口"]}, "task": {"type": "worker", "index": 0}}'

从节点（index=1）的环境变量：

export TF_CONFIG='{"cluster": {"worker": ["主节点IP:端口", "从节点IP:端口"]}, "task": {"type": "worker", "index": 1}}'

然后把策略换成MultiWorkerMirroredStrategy：

strategy = tf.distribute.MultiWorkerMirroredStrategy()

其余代码和单机场景一致。

避坑提示 #

• 确保你的TensorFlow版本是2.x及以上，旧版本的Estimator分布式支持存在不少兼容性问题。
• 如果你的Keras模型用了自定义层，要确保层的逻辑能兼容Estimator的分布式模式（比如不要在层里使用全局变量）。
• 训练时如果出现显存不足，要适当减小单GPU的批量大小，再乘以GPU数量得到总批量。

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