如何在 pytorch 和 tensorflow 中使用张量核？

Question

我正在使用带有张量核心的 Nvidia RTX GPU，我想确保 pytorch/tensorflow 正在使用它的张量核心。 我在几篇文章中注意到张量核心用于处理 float16，默认情况下 pytorch/tensorflow 使用 float32。 他们引入了一些执行“混合精度和分布式训练”的库。 这是一个有点旧的答案。 我想知道 pytorch 或 tensorflow GPU 现在是否支持开箱即用的张量核心处理。

Answer 1

Mixed Precision在两个库中都可用。

对于 pytorch，它是torch.cuda.amp ，自动混合精度 PACKAGE。

https://pytorch.org/docs/stable/amp.html

https://pytorch.org/docs/stable/notes/amp_examples.html 。

Tensorflow 在这里， https://www.tensorflow.org/guide/mixed_precision 。

Answer 2

本页是使用 apex.amp（自动混合精度）的指南，该工具仅在 Python 的 3 行中启用 Tensor Core 加速训练。

您还可以在此处查看顶点 API 的快速入门