feat: track number of batches

sehoffmann · May 10, 2024 · 78755d7 · 78755d7
1 parent a9ce5ac
commit 78755d7
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/dmlcloud/stage.py b/dmlcloud/stage.py
@@ -300,6 +300,8 @@ def train_epoch(self):
             loss = self.train_step(batch)
             self.optimize(loss)
             self.track_reduce(self.loss_metric_name(), loss)
+            self.track_reduce('misc/total_train_batches', torch.tensor(1), reduction=Reduction.SUM, prefixed=False)
+            self.track_reduce('misc/worker_train_batches', torch.tensor(1), reduction=Reduction.SUM, reduce_globally=False, prefixed=False)
 
         for name, scheduler in self.pipeline.schedulers.items():
             self.track(f'misc/lr_{name}', scheduler.get_last_lr()[0], prefixed=False)
@@ -314,6 +316,9 @@ def val_epoch(self):
         for batch in self.val_dataset():
             loss = self.val_step(batch)
             self.track_reduce('loss', loss)
+            self.track_reduce('misc/total_val_batches', torch.tensor(1), reduction=Reduction.SUM, prefixed=False)
+            self.track_reduce('misc/worker_val_batches', torch.tensor(1), reduction=Reduction.SUM, reduce_globally=False, prefixed=False)
+
 
     def table_columns(self):
         columns = super().table_columns()