如何选择在 PyTorch 中使用 float32 矩阵乘法的内部精度？

Question

PyTorch 1.12 将默认的 fp32 数学更改为“最高精度”，并引入了 torch.set_float32_matmul_ precision API，允许用户指定

medium

、

high

和

highest

中的哪个精度用于float32 矩阵乘法的内部精度。

从文档中，我读到选择较低的精度“可能会显着提高性能，并且在某些程序中，精度损失的影响可以忽略不计”。

1。如何确定我的程序是否会从设置较低的精度中受益？这纯粹是经验主义吗？

同样，在使用 PyTorch Lightning 进行训练时，我收到以下警告：

You are using a CUDA device ('NVIDIA A100-SXM4-40GB') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')`

这似乎回答了1。（即，当你的GPU有张量核心时，使用较低的精度），但没有建议使用两个较低精度中的哪一个。

2。如何确定使用哪个较低精度（“高”或“中”）？ this纯粹是经验主义吗？建议的方法是什么？

谢谢！