Mixed precision training

Hi all,

I’m training a model in double precision. However, I notice a significant slowdown compare to single precision counterpart. Therefore, I would like to train part of the model with single precision and the other part in double precision. In my case, the vision backbone (ResNet-34) could be trained with single precision while the rest of the model need to be trained in double precision.

So I’m wondering if there’s any viable way to implement such functionality? Will such mixed precision training effectively improve training speed?

(I know there’s Apex mixed precision training to train part of model in half-precision. Would that framework also work for double precision?)