Is there a way to train in the fp16 mode? about lama

[windj007]

Hi!

Our pipeline supports mixed precision training - just by adding trainer.kwargs.precision=16 +trainer.kwargs.amp_backend=native +trainer.kwargs.amp_level=O1. All arguments in the trainer.kwargs namespace go to pytorch_lightning Trainer - so please refer to their docs for more info.

However, as fft is executed via a function (not a module) pytorch cannot add hooks for type conversion. Thus, it seems that there is no built-in way to enable mixed-precision training with FFT - you have to manually wrap these parts with x.float() and ffted.half().

By the way, we tried to train a Lama-Regular (fully-conv resnet, no FFC) in FP16 and it consistently diverged. However, we did that in the beginning of the project and did not re-check it with final models and a newer Pytorch.

Personally, I often face lack of stability with FP16 in multi-objective training - not only in inpainting. It usually works well with a single loss, but when I try to mix e.g. segmentation + localization or reconstruction + discriminator - loss explodes to nan after a few hundred iterations.

I'd be happy to hear from you if FP16 works finally :)

[ternaus]

I was not able to make it work in FP16 successfully.

But disabling validation and saved a lot of memory on GPU 0 and it is good enough for now :)