Added mixed precision training and removed torch.cuda.empty_cache() d…

…ue to performance drop.

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

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"

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(

luxonis_train/models/luxonis_lightning.py

@@ -10,7 +10,6 @@
 from lightning.pytorch.utilities import rank_zero_only  # type: ignore
 from luxonis_ml.data import LuxonisDataset
 from torch import Size, Tensor, nn
-from torch.amp import autocast
 
 import luxonis_train
 from luxonis_train.attached_modules import (
@@ -394,8 +393,7 @@ def forward(
                 else:
                     node_inputs.append({"features": [inputs[pred]]})
 
-            with autocast(device_type=self.device.type):
-                outputs = node.run(node_inputs)
+            outputs = node.run(node_inputs)
 
             computed[node_name] = outputs
 
@@ -445,8 +443,6 @@ def forward(
             if node_name in self.outputs
         }
 
-        torch.cuda.empty_cache()
-
         return LuxonisOutput(
             outputs=outputs_dict, losses=losses, visualizations=visualizations
         )