我应该使用 fp16 张量来利用张量核心吗？

Question

我正在尝试在AWS上使用V100机器。官方网站说它为机器学习应用程序提供了超过 100TFLOPS，但它似乎没有比我用过的 1080ti 快。

最近发现一篇文章说tensor core运行在fp16张量上。

那么，我应该指定张量的 dtype 吗，比如

tensor_a = torch.tensor([1.,2.,3.], dtype=torch.float16)

而不是：

tensor_a = torch.tensor([1.,2.,3])

？

Answer 1

是的，你应该尝试这个。如果您正在进行推理，您可以手动创建/将张量转换为 fp16，并且您应该会看到显着的加速。如果你想改进训练，你可以使用 torch 的 AMP（自动混合精度），如果认为有帮助的话，它会自动将 fp16 应用于操作。

考虑尝试 fp16 和 bf16。根据我的经验，bf16 更好，但对于您的型号可能会有所不同。请务必测试模型的准确性，以确保它不会降低或变得不稳定。