PyTorch或TensorFlow在训练和运行Tranformer模型时哪个更有效呢?
近日,Huggingface(AI初创企业)的团队成员通过一篇发表在Medium上的博文告诉我们:两者差不多!
Facebook 的首席Ai科学家Yan Lecun大神在社交媒体上对此结果进行了转发点赞。
自TensorFlow发布实施以来,Huggingface团队就致力于模型的产品化,并使它们在TPU上可用,从而逐步提高性能。
这篇文章比较了模型在几种环境下的性能,研究人员在PyTorch(1.3.0)和TensorFlow(2.0)的CPU和GPU上比较它们的推理性能。
以下是博文具体内容:
测试结果
具体结果显示在此下面的Google Spreadsheet中,可点击查看。
https://docs.google.com/spreadsheets/d/1sryqufw2D0XlUH4sq3e9Wnxu5EAQkaohzrJbd5HdQ_w/edit#gid=0
下表中显示了平均结果。结果将在讨论部分中详细介绍。
平均推理时间
电子表格中的N/A项表示内存不足错误或不适当的序列长度。Transformer-XL没有TorchScript结果,因为它目前不能被TorchScript序列化。
在大多数情况下,TensorFlow和PyTorch模型在GPU和CPU上都能获得非常相似的结果。
下面是关于结果的简短讨论,既可以作为PyTorch和TensorFlow之间的比较,也可以作为模型之间的比较。
推理时间测量
在将模型投入生产时,推理时间是一个重要的度量指标。为了评估我们的模型的推理时间,我们将它们与不同的批量大小和不同的序列长度进行比较。我们将合理的批大小[1、2、4、8]与序列长度[8、64、128、256、512、1024]进行比较。批处理大小仍然很小,因为我们只考虑推理设置。BERT和其他类似的模型最大序列长度为512或256(对于CTRL),因此不会测量最后的序列长度。
我们在两种不同的环境中测试结果:
- 在CPU上,使用GCP n1-standard-32,它有32个vcpu和120GB的RAM。CPU型号是Intel Xeon @ 2.3GHz。
- 在GPU上,使用一个定制的GCP机器,它有12个vcpu, 40GB的RAM和一个单独的V100 GPU (16GB的VRAM)。
实验详情和最佳实践
为了使性能最大化,需要进行进一步的优化:
- 我们测量CPU推理的Intel Xeon CPU带有AVX和AVX2扩展。由于TensorFlow需要从源代码编译来利用这些扩展,所以我们这样做。
- 我们通过tf.function并预先对模型进行跟踪,来确保没有使用TensorFlow的eager模式。
- 我们比较了使用和不使用依赖于库的工具的推论:使用PyTorch的TorchScript和使用gpu的TensorFlow的XLA(自动集群)。下面详细介绍这两个工具。
- 我们使用本地Python模块timeit来测量推理时间。我们使用repeat=30和number=3来运行每个实验。然后我们对30个值求平均值,以获得预期的平均推理时间。平均超过30个值会得到非常稳定的结果。
- 我们不使用TFX这样的生产环境,而是测量模型的可调用方法:用于PyTorch的nn.Module.forward和用于TensorFlow的tf.keras.layers.Layer.call
- 我们谨慎地为TensorFlow和PyTorch使用适当的CUDA版本。
讨论
PyTorch和TensorFlow
在大多数情况下,这两个库都能获得相似的结果,与PyTorch相比,TensorFlow在CPU上通常要慢一些,但在GPU上要快一些:
- 在所有模型中,在CPU上,PyTorch的平均推理时间为0.748s,而TensorFlow的平均推理时间为0.823s。
- 在所有模型中,在GPU上,PyTorch的平均推理时间为0.046s,而TensorFlow的平均推理时间为0.043s。
这些结果通过平均结果来比较所有模型的推断时间。因此,输入的大小越大,对最终结果的影响就越大。当输入大小过大时,PyTorch将耗尽内存;当求平均值时,这些结果将从所有度量中删除,因为这样会使结果偏向PyTorch。
PyTorch模型比TensorFlow模型更早耗尽内存:除了Distilled模型外,除了Distilled模型之外,当输入大小达到8的批处理大小和1024的序列长度时,PyTorch就会耗尽内存。
TorchScript
TorchScript是PyTorch创建可序列化模型的方法,该模型可以在不同的运行时运行,而无需Python依赖项,例如C ++环境。我们的测试是通过在Python中跟踪模型并在同一环境中重新使用该跟踪模型来完成的。我们通过预先执行前向遍历来确保在测量其推理之前跟踪模型。
免责声明:虽然在Python环境中,创建TorchScript似乎并不是为了加速,但是我们的结果表明,使用TorchScript跟踪模型可以产生性能改进。
TorchScript似乎非常依赖于模型和输入大小(批量大小*序列长度);例如,使用TorchScript可以在XLNet上产生永久的性能提升,而在XLM上使用它可能会有问题,因为在XLM上,它可以在较小的输入大小下提高性能,但在较大的输入大小下降低性能。
平均而言,使用TorchScript跟踪模型的推理要比使用相同的PyTorch非跟踪模型的推理快20%。
XLA
XLA是可以加速TensorFlow模型的线性代数编译器。我们仅在基于TensorFlow的自动聚类功能的GPU上使用它,该功能可编译一些模型的子图。
结果是提高了速度和内存使用率:启用XLA后,大多数内部基准测试的运行速度提高了约1.15倍。
启用XLA后,我们所有模型的性能都会提高。在某些极端情况下,我们的推理时间减少了70%,特别是在较小的输入大小下。
模型及其简化版本
Distilled模型在该测试中表现出色,因为它非常易于基准测试。这两种工程模型(DistilBERT和DistilGPT-2)与教师模型相比,其推理时间减少了一半。
下一步做什么?
对模型进行基准测试只是我们提高性能之路的第一步。我们认为,在比较模型的当前状态时,尤其是在查看PyTorch和TensorFlow之间的差异时,这篇介绍性文章可能会对您有所帮助。当我们深入研究的生产方面时Transformers,势必会进行面向性能的改进。