节约GPU显存和加速神经网络
安装方式:
pip uninstall apex
git clone https://www.github.com/nvidia/apex
cd apex
python setup.py install
快速应用
from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()