Add mix precision training (#172)

* Add mix precision training

* Disable mixed precision for cpu

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

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

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:

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:

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}
 
 

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

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
 
 

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__":