Mixed precision training

VisCy supports automatic mixed precision (AMP) when training on modern CUDA devices. This feature can be turned on by setting trainer.precision to "16-mixed".

While AMP reduces VRAM usage and can speed up training when the task is GPU-bound, it may have issues with numerical stability and slower convergence.

For example, when training a 3D-LUNeXt model, some voxels can overflow (left: AMP, right: 32-bit):

The training also appears to converge slower (grey: AMP, blue: 32-bit):

These issues do not always occur. For example when training the same model on a different dataset, there was no overflowing observed. In general AMP is a good starting place to improve GPU performance, keeping in mind these common issues when using a lower-precision representation.