TF-Replicator, GPipe, Mesh-Tensorflow 三个库对比

使用场景

从使用场景来看，

TF-Replicator 主要侧重于 data parallelism 的库。虽然也能支持 model parallelism，是基于 parameter server 实现的，用的人不多。这点跟 TF 的 Distribution Strategies 是一样的，但是 tf-replicator 又略有不同：

TF-Replicator 屏蔽 CPU, GPU, TPU 的区别，同一份代码就能在不同的平台上训练，TF Distribution Strategies 针对 GPU, TPU 等不同场景要用不同的 Strategy
TF-Replicator 可以比较细致的控制训练过程， Stragety 比较高层，只能和 Estimator 配合使用，无法定制训练过程

GPipe 是一个纯 model-parallelism 的库，把一个 L 层的网络，切分成 K 个 composite layers, 每个 composite layer 在单独的 TPU core 上，K 个 core 只能顺序执行，但是 GPipe 引入了 pipeline parallelism 来缓解这个顺序执行的性能问题，把 mini-batch 细分为多个更小的 macro-batch, 提高并行程度。最后它还用 recomputation 这个简单有效的技巧来降低内存，进一步允许训练更大的模型。

Mesh-TensorFlow 是一个 DSL 语言，既支持 data parallelism 又支持 parallelism, 目前只支持少量的 Language Model, 没有实现卷积操作 conv，暂时无法用于 CNN 网络。它的特性适合 SuperComputer, 例如 TPU Pod, 简直是为 2D 互联的 TPU 芯片量身定制的，GPU 集群应该不是它的使用场景。

侵入性

上述 3 个库都需要开发者修改代码，具有侵入性，论侵入性的程度：

Mesh-TensorFlow 最严重，需要你改写模型代码和训练过程，且控制的粒度很细，写起来比较繁琐；
GPipe 也需要重写模型代码，不需要修改训练过程，侵入性略低于 Mesh-TensorFlow；
TF-Replicator 不需要改模型代码，只需要少量几行代码用于修改训练过程。

代码的完善程度

从开源代码的完善程度看，

Mesh-TensorFlow 只支持 Language Model, 不支持 CNN, github 文档里明确说了暂时没有实现 conv op
GPipe 目前也只支持 Language Model 这类序列模型，不支持 CNN，github 上的代码只支持 recurrent model, 但是原理上它是很通用的，可以支持 CNN，论文里也在 ImageNet 上训练过了
TF-Replicator 暂时没有开源，只放出了论文，暂时只能用 TF 自带的 Distribution Strategy 代替。

看起来直接使用它们不太现实，只能把它们背后的技术点抽取出来，在自己代码中实现。Mesh-TensowFlow 背后的主要原理应该是 Spatial partition + Graph partition, GPipe 则是 Layer partition + Pipeline parallelism + Recomputation.

最友好的情况，应该是编译器（XLA, TVM 等）底层实现这些优化技巧，在上层库以 annotation 形式提供给用户使用。例如编译器底层实现 recomputation, 用户在 Python 代码里使用 @recomputation，就可以启用 gradient checkpointing, 这样是最方便的。现在 XLA, TVM, MLIR 等项目有在朝着这个方向努力，但都在早期，很不成熟，目前看手工实现这些优化技巧是不可避免了。

如何实现超大模型 (Giant Model)

现在的模型越来越大，提高单个模型的容量迫在眉睫。比如 NAS 这个领域，现在都流行 one-shot 架构，成千上万个子网络在一个 SuperNetwork 里，单卡内存没有可能放得下。FBNet, DARTS,ProxylessNAS, SNAS 都用了一些压缩信息的技巧，虽然能训练出来，但是丢失了很多信息，它们在 ImageNet 上的 top-1 准确率都在 73% 左右徘徊，比目前的 STOA AmoebaNet 整整低了 10 个百分点。

要提高单个模型容量，看起来需要把 partition + pipeline parallelism + recomputation + op fusion 等等技巧全部用起来。

PS: 前几天有篇新论文，是一个局部小优化，Single-Path NAS: Designing Hardware-Efficient ConvNets in less than 4 Hours, https://arxiv.org/abs/1904.02877, 用一个 7x7 的大卷积，masking 后可以用来代表 3x3, 5x5, 7x7 的卷积，这个方法能大幅省 feature maps 数量，三个变成一个，每层都如此，内存指数级下降。两个月前看 weight sharing 相关论文也想到这个点子，没想到这么快已经有人做出来了。

写下你的评论...

你好，想请教一下大佬，有没有可能通过修改 TensorFlow 自带的 Distribution Strategies 来实现更细粒度的控制？

TF-Replicator 允许你自定义训练过程，与其修改 strategy，不如用 tf-replicator

不过关于 TF-Replicator 还没有文档，连 TF2.0.0-Alpha0 也只是个临时的版本，只能再等等了~

TF-Replicator 没有开源，它主要是 data parallelism，功能跟 distribute strategy 重合

有趣，VGG 把 7*7 的卷积改成三层 3*3 的，这篇又改回去了

主要是为了节省 feature maps 的个数

嗯，VGG 是为了省参数，这篇是为了省显存

请问文末说的 weight sharing 的论文是哪篇啊？

就是文末的那个 arXiv 论文

全文完

本文由简悦 SimpRead 优化，用以提升阅读体验

使用了全新的简悦词法分析引擎^beta，点击查看详细说明

使用场景侵入性代码的完善程度如何实现超大模型 (Giant Model)