TensorFlow 2.x 分布式训练必须使用 tf.distribute.Strategy,否则即使多卡或多机也仅单设备运行;所有模型、优化器、数据构建须在 strategy.scope() 内完成,且数据需用 strategy.experimental_distribute_dataset() 包装或由 Keras 自动适配。

TensorFlow 2.x 分布式训练必须用 tf.distribute.Strategy
不用 tf.distribute.Strategy,就不是真正的分布式训练——哪怕你开了多卡或多机,模型仍只在单设备上跑。TensorFlow 2.x 移除了旧的 tf.estimator 分布式 API 和手动图构建方式,所有标准分布式路径都收口到这个抽象层。它不关心你是单机多 GPU、多机多 GPU 还是 TPU,统一通过策略对象封装设备分配、变量创建、梯度同步逻辑。
常见错误现象:VisibleDeviceList 设置了但 GPU 使用率始终为 0;tf.config.list_physical_devices("GPU") 能看到多卡,但 model.fit() 只跑在 /GPU:0 上——本质都是没把模型、数据、优化器的构建放进策略作用域里。
- 必须在创建模型、优化器、指标前调用
strategy.scope(),否则变量不会被复制到各设备 -
tf.keras.Model实例不能在scope()外编译(.compile()),否则会报ValueError: Variable not in scope - 数据必须用
strategy.experimental_distribute_dataset()包装(或直接用tf.data.Dataset.batch().prefetch()后传给model.fit(),Keras 会自动适配)
单机多 GPU 最简可行配置:用 tf.distribute.MirroredStrategy
这是最常用也最容易踩坑的场景。MirroredStrategy 在每个 GPU 上镜像一份模型副本,前向/反向计算并行,梯度通过 NCCL(Linux)或 NVIDIA Collective Communications Library)同步。它不要求修改模型结构,但对数据批大小(batch size)有隐含要求:全局 batch size = 单卡 batch size × GPU 数量。
典型错误:batch_size=32 直接喂给 4 卡训练,结果每卡只收到 8 个样本,小批量导致 BN 层统计失效、梯度不稳定。正确做法是显式放大本地 batch size。
立即学习“Python免费学习笔记(深入)”;
- 初始化策略时加
cross_device_ops=tf.distribute.NcclAllReduce()(Linux)或tf.distribute.ReductionToOneDevice()(macOS/Windows,性能差,仅调试用) - 模型输入层的
batch_size参数(如有)应设为每卡值,例如tf.keras.Input(shape=(224,224,3), batch_size=16)表示每卡处理 16 样本 - 验证集评估默认不启用分布式,如需多卡验证,得手动写
@tf.function+strategy.run(),否则会退化成单卡预测
自定义训练循环中必须用 strategy.run() 和 strategy.reduce()
一旦脱离 model.fit(),比如你要做梯度裁剪、混合精度、多任务 loss 加权,就必须手写训练步。此时不能直接调 model(x) 或 optimizer.apply_gradients(),因为这些操作不在设备上下文中。
错误写法:loss = model(x, training=True) → 报错 Op is not supported on this device;正确写法是把整个前向+反向逻辑封装进函数,再交由 strategy.run() 分发。
- 每个设备上的 loss 是局部的,需用
strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_loss, axis=None)汇总为标量 -
strategy.run()返回的是PerReplica对象,不能直接 print 或转 numpy,要先strategy.experimental_local_results()或strategy.reduce() - 混合精度(
tf.keras.mixed_precision.Policy("mixed_float16"))必须在strategy.scope()内设置,且优化器要用tf.keras.mixed_precision.LossScaleOptimizer包装
多机训练前先确认网络与 TF_CONFIG 环境变量
多机不是加几行代码就能跑通的事。TensorFlow 依赖 gRPC 通信,默认端口 8470(TPU)或 8500(CPU/GPU),若防火墙拦截、端口被占、主机名无法解析,进程会卡在 Waiting for other workers 或直接 timeout。
TF_CONFIG 是唯一调度依据,格式严格:必须是 JSON 字符串,cluster 下分 worker、ps、chief 等角色,task 中指定 index 和 type。少一个引号或逗号都会让 tf.distribute.MultiWorkerMirroredStrategy 初始化失败。
- 所有机器必须使用完全相同的代码、Python 版本、TensorFlow 版本、CUDA/cuDNN 版本
- 推荐用
nccl而非ring通信后端,后者在跨机器时易出同步问题 - 首次部署建议先关掉 eager execution(
tf.compat.v1.disable_eager_execution()),避免某些 trace 错误干扰定位
真正麻烦的从来不是写几行 strategy 代码,而是网络连通性、NCCL 共享内存限制(/dev/shm 太小会导致 all-reduce hang)、以及不同机器间数据路径不一致导致的 FileNotFoundError ——这些都不会在单机测试中暴露。


















