Skip to content

Mixed precision training

Ziwen Liu edited this page Jun 23, 2025 · 2 revisions

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 UNeXt2 model, some voxels can overflow (left: AMP, right: 32-bit):

image

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

image

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.

Clone this wiki locally