Mixed Precision Training (#113)

luxonis · Oct 31, 2024 · 1d83a2c · 1d83a2c
1 parent 4fe0570
commit 1d83a2c
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 11 deletions.
diff --git a/luxonis_train/attached_modules/losses/adaptive_detection_loss.py b/luxonis_train/attached_modules/losses/adaptive_detection_loss.py
@@ -270,8 +270,11 @@ def forward(
             self.alpha * pred_score.pow(self.gamma) * (1 - label)
             + target_score * label
         )
-        ce_loss = F.binary_cross_entropy(
-            pred_score.float(), target_score.float(), reduction="none"
-        )
+        with torch.amp.autocast(
+            device_type=pred_score.device.type, enabled=False
+        ):
+            ce_loss = F.binary_cross_entropy(
+                pred_score.float(), target_score.float(), reduction="none"
+            )
         loss = (ce_loss * weight).sum()
         return loss
diff --git a/luxonis_train/config/config.py b/luxonis_train/config/config.py
@@ -340,6 +340,7 @@ class TrainerConfig(BaseModelExtraForbid):
     preprocessing: PreprocessingConfig = PreprocessingConfig()
     use_rich_progress_bar: bool = True
 
+    precision: Literal["16-mixed", "32"] = "32"
     accelerator: Literal["auto", "cpu", "gpu", "tpu"] = "auto"
     devices: int | list[int] | str = "auto"
     strategy: Literal["auto", "ddp"] = "auto"

diff --git a/luxonis_train/core/core.py b/luxonis_train/core/core.py
@@ -108,6 +108,7 @@ def __init__(
             callbacks=LuxonisRichProgressBar()
             if self.cfg.trainer.use_rich_progress_bar
             else LuxonisTQDMProgressBar(),
+            precision=self.cfg.trainer.precision,
         )
 
         self.train_augmentations = Augmentations(

diff --git a/luxonis_train/models/luxonis_lightning.py b/luxonis_train/models/luxonis_lightning.py
@@ -392,9 +392,13 @@ def forward(
                     node_inputs.append(computed[pred])
                 else:
                     node_inputs.append({"features": [inputs[pred]]})
+
             outputs = node.run(node_inputs)
+
             computed[node_name] = outputs
 
+            del node_inputs
+
             if (
                 compute_loss
                 and node_name in self.losses
@@ -420,20 +424,15 @@ def forward(
                     node_name
                 ].items():
                     viz = combine_visualizations(
-                        visualizer.run(
-                            images,
-                            images,
-                            outputs,
-                            labels,
-                        ),
+                        visualizer.run(images, images, outputs, labels),
                     )
                     visualizations[node_name][viz_name] = viz
 
             for computed_name in list(computed.keys()):
                 if computed_name in self.outputs:
                     continue
-                for node_name in unprocessed:
-                    if computed_name in self.graph[node_name]:
+                for unprocessed_name in unprocessed:
+                    if computed_name in self.graph[unprocessed_name]:
                         break
                 else:
                     del computed[computed_name]