MIR.utils.precision

Mixed-precision helpers: one precision switch for training and inference.

MIR models run in float32 (default), float16 or bfloat16 mixed precision (torch.autocast) and in pure float16 inference (model.half()). Losses and regularizers always compute in float32 and return a float32 value; spatial transformers, integrators and Fourier kernels compute their geometry in float32 and return the input dtype.

Example:

from MIR.utils.precision import autocast, grad_scaler

scaler = grad_scaler(precision)                    # active only for fp16 on CUDA
with autocast(precision):
    flow = model((mov, fix))
    loss = sim(warped, fix) + reg(flow)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

Functions

autocast([precision, device_type])	Context manager for the requested precision (a no-op for "fp32").
check_precision(precision[, device_type])	Raise a clear error when precision is not usable on device_type.
grad_scaler([precision, device_type])	Gradient scaler for the requested precision.
precision_dtype(precision)	Return the torch dtype for "fp32", "fp16" or "bf16".