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

Posted 2023-02-16

【中文标题】如何在 pytorch 和 tensorflow 中使用张量核心？

【英文标题】：How to use tensor cores in pytorch and tensorflow?

【发布时间】：2021-12-28 23:30:01

【问题描述】：

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

【问题讨论】：

在 pytorch 中，您必须使用模型和输入中的 .cuda() 方法将模型移动到 gpu。然后你可以在nvidia软件中查看使用情况nvidia-smi

【参考方案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。