迁移tf.contrib.learn：配置含独立评估进程的分布式训练

咱们一步步来调整代码，适配你的需求——两台GPU机器做分布式训练，同时用一个CPU进程持续跑小型验证集评估：

1. 配置分布式集群结构 #

首先得明确集群里各节点的角色，区分GPU训练节点和CPU评估节点。假设你有两台GPU机器（命名为worker0、worker1），一台CPU评估机器（命名为evaluator），集群配置可以这么写：

cluster_spec = tf.train.ClusterSpec({
    "worker": ["worker0:2222", "worker1:2222"],
    "evaluator": ["evaluator:2223"]
})

然后给每个进程指定对应的角色和任务ID：

• GPU训练节点：task_type="worker"，task_id分别设为0和1
• CPU评估节点：task_type="evaluator"，task_id=0

2. 配置Estimator运行环境 #

为训练和评估节点分别创建RunConfig，核心是所有节点共用同一个共享模型目录（必须是NFS、HDFS这类分布式文件系统，不然评估节点读不到训练生成的checkpoint）：

# GPU训练节点的配置
if task_type == "worker":
    config = tf.estimator.RunConfig(
        cluster=cluster_spec,
        task_type=task_type,
        task_id=task_id,
        model_dir="/path/to/shared/model_dir",
        save_checkpoints_steps=1000,  # 按需调整checkpoint保存频率
        log_step_count_steps=100
    )
# CPU评估节点的配置
elif task_type == "evaluator":
    config = tf.estimator.RunConfig(
        cluster=cluster_spec,
        task_type=task_type,
        task_id=task_id,
        model_dir="/path/to/shared/model_dir"  # 和训练节点共用同一目录
    )

3. 拆分训练/评估输入函数 #

分别定义训练和验证集的输入逻辑，评估集要关闭shuffle：

def train_input_fn():
    # 替换成你的训练数据加载逻辑
    dataset = tf.data.Dataset.from_tensor_slices(...)
    dataset = dataset.shuffle(10000).batch(64).repeat()
    return dataset

def eval_input_fn():
    # 替换成你的小型验证集加载逻辑
    dataset = tf.data.Dataset.from_tensor_slices(...)
    dataset = dataset.batch(64).repeat(1)  # 验证集只跑一遍
    return dataset

4. 配置TrainSpec和EvalSpec #

这两个是train_and_evaluate的核心参数，重点设置评估的持续触发规则：

# 训练任务配置
train_spec = tf.estimator.TrainSpec(
    input_fn=train_input_fn,
    max_steps=100000,  # 总训练步数，按需调整
    hooks=[tf.train.StopAtStepHook(last_step=100000)]  # 可选，添加训练钩子
)

# 持续评估配置
eval_spec = tf.estimator.EvalSpec(
    input_fn=eval_input_fn,
    steps=None,  # 评估整个验证集
    start_delay_secs=60,  # 先等60秒让训练生成第一个checkpoint
    throttle_secs=300,  # 每300秒触发一次评估，按需调整
    hooks=[tf.train.LoggingTensorHook({"loss": "loss"}, every_n_iter=10)]  # 可选，添加评估日志钩子
)

5. 启动各节点任务 #

最后在对应机器上运行代码：

# 替换成你自己的模型Estimator
estimator = tf.estimator.DNNClassifier(
    hidden_units=[256, 128],
    feature_columns=feature_columns,
    n_classes=10,
    config=config
)

# 启动训练/评估流程
tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

• 在两台GPU机器上运行设置task_type="worker"的代码
• 在CPU机器上运行设置task_type="evaluator"的代码，还可以通过export CUDA_VISIBLE_DEVICES=""强制让评估进程只用CPU

关键注意事项 #

• 必须用分布式共享目录存储模型，本地目录会导致评估节点无法读取训练生成的checkpoint
• CPU评估节点的evaluator是TF 1.8支持的自定义角色，不需要加入训练集群的worker/ps组
• 训练节点要确保TF能正确识别GPU资源，避免和评估进程冲突

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