diff --git a/README.md b/README.md index 4a06916..b54bb4c 100644 --- a/README.md +++ b/README.md @@ -265,6 +265,12 @@ Options: --device TEXT Device name (e.g. "cuda", "cuda:0", "cpu"). Defaults to "cuda" --fp16 Use fp16 precision for much faster inference. + --bf16 Run the ViT encoder in bf16 (weights and + compute) and keep the neck, heads and refiner in + fp32, i.e. the training-time mixed precision + policy. Roughly halves the weight memory. v2/v3 + only, requires a CUDA GPU with native bf16 + support (Ampere or newer). --resize INTEGER Resize the image(s) & output maps to a specific size. Defaults to None (no resizing). --resolution_level INTEGER An integer [0-9] for the resolution level for diff --git a/moge/model/v2.py b/moge/model/v2.py index c5aafdc..4157a6f 100644 --- a/moge/model/v2.py +++ b/moge/model/v2.py @@ -69,7 +69,10 @@ def device(self) -> torch.device: @property def dtype(self) -> torch.dtype: - return next(self.parameters()).dtype + # The dtype inputs are expected in. Taken from the neck rather than the first parameter, because under + # `enable_mixed_precision(..., cast_encoder_weights=True)` the encoder is stored in reduced precision while + # the neck and heads (and therefore the inputs and outputs) stay in fp32. + return next(self.neck.parameters()).dtype @property def onnx_compatible_mode(self) -> bool: @@ -138,14 +141,22 @@ def enable_gradient_checkpointing(self): if hasattr(self, head): getattr(self, head).enable_gradient_checkpointing() - def enable_mixed_precision(self, dtype: torch.dtype = torch.bfloat16): + def enable_mixed_precision(self, dtype: torch.dtype = torch.bfloat16, cast_encoder_weights: bool = False): """Enable fine-grained mixed precision: run the encoder in `dtype`, keep the neck and heads in fp32. Calling this repeatedly replaces the previous wrapping rather than stacking it. + + If `cast_encoder_weights` is True, the encoder weights are also stored in `dtype`. Autocast casts them + to `dtype` on the fly anyway, caching the copies for the duration of the forward pass, so this does not + change the precision of the matmuls but roughly halves the encoder's weight memory and avoids holding + the weights twice. Intended for inference. """ for handle in getattr(self, '_autocast_handles', []): handle.remove() + if cast_encoder_weights: + self.encoder.to(dtype) + module_dtype_map = [ (self.encoder, dtype), (self.neck, torch.float32), diff --git a/moge/scripts/infer.py b/moge/scripts/infer.py index ef25bd4..c68024d 100644 --- a/moge/scripts/infer.py +++ b/moge/scripts/infer.py @@ -20,6 +20,8 @@ @click.option('--version', 'model_version', type=click.Choice(['v1', 'v2', 'v3']), default='v3', help='Model version. Defaults to "v3"') @click.option('--device', 'device_name', type=str, default='cuda', help='Device name (e.g. "cuda", "cuda:0", "cpu"). Defaults to "cuda"') @click.option('--fp16', 'use_fp16', is_flag=True, help='Use fp16 precision for much faster inference.') +@click.option('--bf16', 'use_bf16', is_flag=True, help='Run the ViT encoder in bf16 (weights and compute) and keep the neck, heads and refiner in fp32, \ +i.e. the training-time mixed precision policy. Roughly halves the weight memory. v2/v3 only, requires a CUDA GPU with native bf16 support (Ampere or newer).') @click.option('--resize', 'resize_to', type=int, default=None, help='Resize the image(s) & output maps to a specific size. Defaults to None (no resizing).') @click.option('--resolution_level', type=int, default=9, help='An integer [0-9] for the resolution level for inference. \ Higher value means more tokens and the finer details will be captured, but inference can be slower. \ @@ -41,6 +43,7 @@ def main( model_version: str, device_name: str, use_fp16: bool, + use_bf16: bool, resize_to: int, resolution_level: int, num_tokens: int, @@ -86,9 +89,15 @@ def main( if model_version == 'v3': raise click.UsageError('MoGe-3 checkpoints are not released to Huggingface yet. Please provide a local path to the checkpoint.') pretrained_model_name_or_path = default_pretrained_models[model_version] + if use_bf16 and model_version == 'v1': + raise click.UsageError('--bf16 is only supported for v2 and v3.') + if use_bf16 and use_fp16: + raise click.UsageError('--bf16 and --fp16 are mutually exclusive.') model = import_model_class_by_version(model_version).from_pretrained(pretrained_model_name_or_path).to(device).eval() if use_fp16 and model_version != 'v3': model.half() + if use_bf16: + model.enable_mixed_precision(torch.bfloat16, cast_encoder_weights=True) if not any([save_maps_, save_glb_, save_ply_]): warnings.warn('No output format specified. Defaults to saving all. Please use "--maps", "--glb", or "--ply" to specify the output.')