Warmup changes: only warm a few batches; extract to separate method in trainer class (#43)

* apply optimizer every batch, not every epoch; unscale gradients before clipping

* trainer tweaks

* apply optimizer every batch, not every epoch; unscale gradients before clipping

* extract warmup to separate method; switch to warming up set number of batches (user configurable)

* whitespace; num_workers revert

* ruff

* make parallelstrategy, spatial_mesh, ddp_placements attrs of trainer; other small tweaks

* remove deprecated config attrs

* ruff

* get device mesh from ps class attr

* ruff

* missing self. on some ps accesses

* Fix imports and missing self.ps

* rm legacy warmup_epochs

* Move attributes to base class for clarity

* remove warmup_epochs -- not useful to keep support for this

* call cleanup_or_resume trainer method directly

* rm unused vars

---------

Co-authored-by: Patrick Miles <miles30@tioga.llnl.gov>
Co-authored-by: Michael McKinsey <michaelmckinsey1@gmail.com>

Author: 3 people

ScaFFold/cli.py

@@ -140,6 +140,11 @@ def main():
     benchmark_parser.add_argument(
         "--batch-size", type=int, nargs="+", help="Batch sizes for each volume size."
     )
+    benchmark_parser.add_argument(
+        "--warmup-batches",
+        type=int,
+        help="Number of warmup batches to run per rank before training.",
+    )
     benchmark_parser.add_argument(
         "--optimizer",
         type=str,

ScaFFold/configs/benchmark_default.yml

@@ -29,6 +29,6 @@ framework: "torch"                 # The DL framework to train with. Only valid
 checkpoint_dir: "checkpoints"      # Subfolder in which to save training checkpoints.
 loss_freq: 1                       # Number of epochs between logging the overall loss.
 normalize: 1                       # Cateogry search normalization parameter
-warmup_epochs: 1                   # How many warmup epochs before training
+warmup_batches: 5                  # How many warmup batches per rank to run before training.
 dataset_reuse_enforce_commit_id: 0 # Enforce matching commit IDs for dataset reuse.
-target_dice: 0.95
+target_dice: 0.95

ScaFFold/configs/benchmark_testing.yml

@@ -29,6 +29,6 @@ framework: "torch"                 # The DL framework to train with. Only valid
 checkpoint_dir: "checkpoints"      # Subfolder in which to save training checkpoints.
 loss_freq: 1                       # Number of epochs between logging the overall loss.
 normalize: 1                       # Cateogry search normalization parameter
-warmup_epochs: 1                   # How many warmup epochs before training
+warmup_batches: 5                  # How many warmup batches per rank to run before training.
 dataset_reuse_enforce_commit_id: 0 # Enforce matching commit IDs for dataset reuse.
 target_dice: 0.95

ScaFFold/utils/config_utils.py

@@ -66,7 +66,7 @@ def __init__(self, config_dict):
         self.loss_freq = config_dict["loss_freq"]
         self.checkpoint_dir = config_dict["checkpoint_dir"]
         self.normalize = config_dict["normalize"]
-        self.warmup_epochs = config_dict["warmup_epochs"]
+        self.warmup_batches = config_dict.get("warmup_batches")
         self.dataset_reuse_enforce_commit_id = config_dict[
             "dataset_reuse_enforce_commit_id"
         ]