add mixed precision training/matmul32 precision

src/deepforest/conf/config.yaml

@@ -17,6 +17,13 @@ model:
     name: 'weecology/deepforest-tree'
     revision: 'main'
 
+# Trainer precision. Override to 16-mixed for faster training, 32-true for full precision.
+precision:
+
+# On CUDA, setting the matmul precision can provide speed up without affecting model performance
+# Start with 'high' for initial tests.
+matmul_precision: highest
+
 # Specify a label_dict to override model settings.
 # By default, this will be populated from the model
 # checkpoint that is selected in model.name/revision.
@@ -77,6 +84,7 @@ train:
     # preload images to GPU memory for fast training. This depends on GPU size and number of images.
     preload_images: False
 
+
 validation:
     csv_file:
     root_dir:

src/deepforest/conf/schema.py

@@ -133,6 +133,8 @@ class Config:
     devices: int | str = "auto"
     accelerator: str = "auto"
     batch_size: int = 1
+    precision: str | None = None
+    matmul_precision: str = "highest"
 
     architecture: str = "retinanet"
     num_classes: int | None = None

src/deepforest/main.py

@@ -231,6 +231,9 @@ def create_trainer(self, logger=None, callbacks=None, **kwargs):
         else:
             enable_checkpointing = False
 
+        if torch.cuda.is_available():
+            torch.set_float32_matmul_precision(self.config.matmul_precision)
+
         trainer_args = {
             "logger": logger,
             "max_epochs": self.config.train.epochs,
@@ -243,6 +246,8 @@ def create_trainer(self, logger=None, callbacks=None, **kwargs):
             "num_sanity_val_steps": num_sanity_val_steps,
             "default_root_dir": self.config.log_root,
         }
+        if self.config.precision is not None:
+            trainer_args["precision"] = self.config.precision
         # Update with kwargs to allow them to override config
         trainer_args.update(kwargs)