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

Question

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

Answer 1

Mixed Precision在两个库中都可用。

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

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

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

Tensorflow 在这里有它，https://www.tensorflow.org/guide/mixed_ precision。