pyro-ppl/pyro

[FR] Support Automatic Mixed Precision training

Open

#3,316 建立於 2024年1月31日

在 GitHub 查看
 (7 留言) (0 反應) (0 負責人)Python (981 fork)batch import
enhancementhelp wanted

倉庫指標

Star
 (8,211 star)
PR 合併指標
 (平均合併 10天 19小時) (30 天內合併 1 個 PR)

描述

Issue Description

Better support for mixed precision training would be extremely helpful, at least for SVI. I can manually cast data into float16 or bfloat16 but I am unable to leverage PyTorch's automatic mixed precision training. This is because it requires the use of the GradScaler class during the optimization loop to properly scale gradients in a mixed-precision-aware manner. See the documentation for more info: https://pytorch.org/docs/stable/amp.html

It would be nice to have support for using this class within pyro optimizers to allow for amp support.

貢獻者指南