diff --git a/configs/graph_default.yaml b/configs/graph_default.yaml index d5c08b7a6..27b358e3a 100644 --- a/configs/graph_default.yaml +++ b/configs/graph_default.yaml @@ -30,6 +30,7 @@ trainer: steps_per_save: 1000 steps_per_test: 1000 max_num_iterations: 1000000 + mixed_precision: False resume_train: load_dir: # pretrained model directory load_step: # pretrained model's step number diff --git a/configs/graph_instant_ngp.yaml b/configs/graph_instant_ngp.yaml index 935428eec..e8b850c42 100644 --- a/configs/graph_instant_ngp.yaml +++ b/configs/graph_instant_ngp.yaml @@ -29,3 +29,6 @@ optimizers: lr: 1e-2 eps: 1e-7 # need to investigate further for float16 stability scheduler: null + +trainer: + mixed_precision: true # Stabilizes training when using tcnn diff --git a/pyrad/engine/trainer.py b/pyrad/engine/trainer.py index 3bf8ad506..20a065634 100644 --- a/pyrad/engine/trainer.py +++ b/pyrad/engine/trainer.py @@ -22,6 +22,7 @@ from typing import Dict import torch import torch.distributed as dist +from torch.cuda.amp import GradScaler from torch.nn.parallel import DistributedDataParallel as DDP from torchtyping import TensorType @@ -52,6 +53,10 @@ class Trainer: self.local_rank = local_rank self.world_size = world_size self.device = "cpu" if world_size == 0 else f"cuda:{local_rank}" + self.mixed_precision = self.config.trainer.mixed_precision + if self.device == "cpu": + self.mixed_precision = False + logging.warning("Mixed precision is disabled for CPU training.") # dataset variables self.dataset_inputs_train = None self.dataloader_train = None @@ -66,6 +71,8 @@ class Trainer: # visualizer variable self.visualizer_state = viewer_utils.VisualizerState(config.viewer) + self.grad_scaler = GradScaler(enabled=self.mixed_precision) + def setup(self, test_mode=False): """Setup the Trainer by calling other setup functions. @@ -147,6 +154,7 @@ class Trainer: # load the checkpoints for graph and optimizer self.graph.load_graph(loaded_state) self.optimizers.load_optimizers(loaded_state) + self.grad_scaler.load_state_dict(loaded_state["scaler"]) logging.info("done loading checkpoint from %s", load_path) @check_main_thread @@ -165,6 +173,7 @@ class Trainer: "step": step, "model": self.graph.module.state_dict() if hasattr(self.graph, "module") else self.graph.state_dict(), "optimizers": {k: v.state_dict() for (k, v) in self.optimizers.optimizers.items()}, + "scalers": self.grad_scaler.state_dict(), }, ckpt_path, ) @@ -181,11 +190,14 @@ class Trainer: Returns: Dict[str, float]: Dictionary of model losses. """ - _, loss_dict = self.graph.forward(ray_indices, batch=batch) - loss = loss_dict["aggregated_loss"] self.optimizers.zero_grad_all() - loss.backward() - self.optimizers.optimizer_step_all() + with torch.autocast(device_type=ray_indices.device.type, enabled=self.mixed_precision): + _, loss_dict = self.graph.forward(ray_indices, batch=batch) + loss = loss_dict["aggregated_loss"] + self.grad_scaler.scale(loss).backward() + self.optimizers.optimizer_scaler_step_all(self.grad_scaler) + self.grad_scaler.update() + self.optimizers.scheduler_step_all(step) if self.graph.callbacks: for func_ in self.graph.callbacks: diff --git a/pyrad/fields/base.py b/pyrad/fields/base.py index 6ac176e1a..43e11911a 100644 --- a/pyrad/fields/base.py +++ b/pyrad/fields/base.py @@ -91,14 +91,14 @@ class Field(nn.Module): field_outputs = {} for k, value in field_outputs_masked.items(): zeros = torch.zeros( - *valid_mask.shape[:-1], value.shape[-1], dtype=torch.float32, device=valid_mask.device + *valid_mask.shape[:-1], value.shape[-1], dtype=value.dtype, device=valid_mask.device ) if valid_mask.any(): zeros[valid_mask[..., 0]] = value else: zeros[0, :] = value field_outputs[k] = zeros - density = torch.zeros(valid_mask.shape, dtype=torch.float32, device=valid_mask.device) + density = torch.zeros(valid_mask.shape, dtype=density_masked.dtype, device=valid_mask.device) if valid_mask.any(): density[valid_mask[..., 0]] = density_masked else: diff --git a/pyrad/fields/instant_ngp_field.py b/pyrad/fields/instant_ngp_field.py index 6a35858e2..7cee20855 100644 --- a/pyrad/fields/instant_ngp_field.py +++ b/pyrad/fields/instant_ngp_field.py @@ -107,14 +107,12 @@ class TCNNInstantNGPField(Field): """Computes and returns the densities.""" positions = get_normalized_positions(ray_samples.frustums.get_positions(), self.aabb) positions_flat = positions.view(-1, 3) - dtype = positions_flat.dtype - h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1).to(dtype) + h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1) density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1) # Rectifying the density with an exponential is much more stable than a ReLU or # softplus, because it enables high post-activation (float32) density outputs # from smaller internal (float16) parameters. - assert density_before_activation.dtype is torch.float32 density = trunc_exp(density_before_activation) return density, base_mlp_out @@ -123,11 +121,9 @@ class TCNNInstantNGPField(Field): # tcnn requires directions in the range [0,1] directions = get_normalized_directions(ray_samples.frustums.directions) directions_flat = directions.view(-1, 3) - dtype = directions_flat.dtype d = self.direction_encoding(directions_flat) h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) - rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(dtype) - assert rgb.dtype is torch.float32 + rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1) return {FieldHeadNames.RGB: rgb} diff --git a/pyrad/optimizers/optimizers.py b/pyrad/optimizers/optimizers.py index 41eac9a4c..347dcfc62 100644 --- a/pyrad/optimizers/optimizers.py +++ b/pyrad/optimizers/optimizers.py @@ -21,6 +21,7 @@ from typing import Any, Dict, List import numpy as np from hydra.utils import instantiate from omegaconf import DictConfig +from torch.cuda.amp import GradScaler from torch.nn import Parameter from torch.optim.lr_scheduler import LambdaLR @@ -115,7 +116,16 @@ class Optimizers: for _, optimizer in self.optimizers.items(): optimizer.zero_grad() - def optimizer_step_all(self) -> None: + def optimizer_scaler_step_all(self, grad_scaler: GradScaler) -> None: + """Take an optimizer step using a grad scaler. + + Args: + grad_scaler: GradScaler to use + """ + for _, optimizer in self.optimizers.items(): + grad_scaler.step(optimizer) + + def optimizer_step_all(self): """Run step for all optimizers.""" for _, optimizer in self.optimizers.items(): # note that they key is the parameter name diff --git a/pyrad/utils/config.py b/pyrad/utils/config.py index f6ff41676..481969cd1 100644 --- a/pyrad/utils/config.py +++ b/pyrad/utils/config.py @@ -56,6 +56,7 @@ class TrainerConfig: steps_per_save: int = MISSING steps_per_test: int = MISSING max_num_iterations: int = MISSING + mixed_precision: bool = MISSING resume_train: ResumeTrainConfig = MISSING diff --git a/tests/fields/test_fields.py b/tests/fields/test_fields.py index c28b274e3..b20a65155 100644 --- a/tests/fields/test_fields.py +++ b/tests/fields/test_fields.py @@ -3,7 +3,6 @@ Test the fields """ import torch -from pyrad.fields.modules.field_heads import FieldHeadNames from pyrad.fields.instant_ngp_field import TCNNInstantNGPField from pyrad.cameras.rays import Frustums, RaySamples @@ -36,13 +35,7 @@ def test_tcnn_instant_ngp_field(): pixel_area=torch.ones((*directions.shape[:-1], 1), device=device), ) ray_samples = RaySamples(frustums=frustums) - field_outputs = field.forward(ray_samples) - - positions_dtype = positions.dtype - field_outputs_density_dtype = field_outputs[FieldHeadNames.DENSITY].dtype - assert ( - positions_dtype == field_outputs_density_dtype - ), f"The input and output dtypes do not match: {positions_dtype} vs. {field_outputs_density_dtype}" + field.forward(ray_samples) if __name__ == "__main__":