在tensorflow中进行模型并行化的正确方法是什么？

Question

我有多个4GB GPU节点，所以我希望它们并行运行大型模型。我希望只需将具有适当设备范围的层拆分为多个部分，就可以实现模型并行性，但事实证明它不会减少主节点的内存占用（任务0）。 （10个节点配置 - 主站：20g，粉丝：2g，1个节点配置 - 主站：6~7g）

可疑的是渐变不会分发，因为我没有为它们设置正确的设备范围。

我的模型可以在github上找到。（https://github.com/nakosung/tensorflow-wavenet/tree/model_parallel_2）

设备放置日志位于：https://gist.github.com/nakosung/a38d4610fff09992f7e5569f19eefa57

Answer 1

所以好消息是你使用了colocate_gradients_with_ops，这意味着你要确保在放置操作的同一设备上计算渐变。 （https://github.com/nakosung/tensorflow-wavenet/blob/model_parallel_2/train.py#L242）

阅读设备放置日志有点困难，因此我建议使用TensorBoard尝试可视化图形。它具有可视化节点如何放置在设备上的选项。

其次，您可以尝试查看操作的大小如何映射到设备 - 最大的层（最大激活或最大权重）可能不成比例地放置在某些节点上而不是其他节点上。您可以尝试使用https://github.com/tensorflow/tensorflow/blob/6b1d4fd8090d44d20fdadabf06f1a9b178c3d80c/tensorflow/python/tools/graph_metrics.py分析图表，以便更好地了解图表中所需资源的位置。

从长远来看，我们希望尝试自动解决其中一些放置问题，但到目前为止，模型并行性需要谨慎处理。