如何在 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 的快速入門