Cloud ML Engine训练MobileNet V1 SSD时参数服务器与workerCount比例咨询

我刚帮不少开发者调过Cloud ML Engine上TensorFlow Object Detection的分布式训练配置，针对你用MobileNet V1 SSD的场景，给你梳理下关键配置的思路，帮你有效缩短训练时间。

一、关于workerCount的设置 #

• 虽然官方提到workerCount的优化还处于试验阶段，但对于MobileNet V1 SSD这类计算量中等的轻量级检测模型，适当提升workerCount确实能明显加速训练——不过前提是你的数据集输入 pipeline 能跟上节奏，比如已经用了TFRecord格式，开启了并行读取和预取。
• 别盲目堆worker数量：当worker数超过某个阈值后，分布式节点间的通信开销会抵消掉计算收益。建议从workerCount=2开始测试，逐步增加到4、6，观察训练速度的提升幅度，直到速度不再明显增长就停住。

二、parameterServerCount（PS节点数）的配置原则 #

• PS节点的核心作用是保存和同步更新模型参数，它的数量必须和worker数匹配，不然很容易成为训练的瓶颈。
• 给你几个通用的经验法则：
  • 当workerCount ≤ 4时，parameterServerCount=1就足够了，这时参数更新的通信量不大，单个PS节点完全能扛住。
  • 当workerCount在4-8之间时，设置parameterServerCount=2，分摊参数更新的压力，避免单个PS拖慢整个训练流程。
  • 如果workerCount超过8，可以按照workerCount:parameterServerCount = 4:1的比例来配置（比如8个worker对应2个PS，12个worker对应3个PS）。
• 另外要注意PS节点的机器类型，别用太弱小的实例，尤其是当模型参数较多时，弱小的PS会导致参数更新延迟严重。

三、parameterServerType与workerCount的比例关系 #

• parameterServerType指的是PS节点的机器规格（比如n1-standard-2、n1-highmem-4这类），它和workerCount的匹配核心是PS的算力和内存要能支撑所有worker的参数更新请求。
• 针对MobileNet V1 SSD这种参数规模不算特别大的模型：
  • 当workerCount ≤ 4时，PS节点用和worker相同的规格（比如worker用n1-standard-4，PS也用n1-standard-4）就完全够用。
  • 当workerCount在4-8之间，建议把PS升级到更高内存的实例（比如n1-highmem-4），因为此时参数更新的并发请求变多，更高的内存能减少参数交换的延迟。
  • 如果workerCount超过8，建议PS节点的规格比worker高一个档次，同时按照之前的4:1比例设置PS数量，确保参数更新的效率。

四、额外的加速小技巧 #

• 开启混合精度训练：MobileNet V1 SSD完全支持混合精度，在Cloud ML Engine上可以通过设置--config=config.yaml，在yaml配置文件里加上use_mixed_precision: true，能大幅提升训练速度，同时几乎不影响检测精度。
• 优化输入 pipeline：确保你的TFRecord文件是分片存储的，每个分片大小控制在100-200MB左右，同时在输入函数里设置num_parallel_reads=8，开启预取prefetch(tf.data.AUTOTUNE)，避免数据读取成为训练的瓶颈。
• 调整batch size：分布式训练时，每个worker的batch size可以适当调大（比如单worker用32，分布式时每个worker保持32，总batch size就是32*workerCount），只要不超出GPU内存（如果用GPU worker的话），更大的batch能提升计算效率。

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