Timing training with timed decorator

requirements.txt

@@ -1,4 +1,6 @@
 torch
 torchvision
 bs_scheduler
-timed-decorator
+timed-decorator>=1.3.0
+tqdm
+argparse

utils/scheduler.py

@@ -1,4 +1,3 @@
-# pip install git+https://github.com/ancestor-mithril/bs_scheduler.git@master
 from bs_scheduler import IncreaseBSOnPlateau
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import ReduceLROnPlateau

utils/trainer.py

@@ -1,9 +1,9 @@
 import os
 import re
 from datetime import datetime
-from timeit import default_timer
 
 import torch
+from timed_decorator.simple_timed import timed
 from torch.backends import cudnn
 from torch.utils.tensorboard import SummaryWriter
 from torchvision.transforms.v2.functional import hflip, vflip  # noqa: F401
@@ -62,16 +62,13 @@ def get_bs(self):
         return self.train_loader.batch_sampler.batch_size
 
     def run(self):
-        # TODO: Implement return time instead of printing
         epochs = list(range(self.args.epochs))
         total_training_time = 0
         try:
             with tqdm(epochs) as tbar:
                 for epoch in tbar:
-                    start = default_timer()
-                    metrics = self.train()
-                    end = default_timer()
-                    total_training_time += end - start
+                    metrics, elapsed = self.train()
+                    total_training_time += elapsed / 1e+9
 
                     metrics.update(self.val())
 
@@ -84,6 +81,7 @@ def run(self):
             f.write(f'{self.logdir} -> {self.best_metric}\n')
         print("Best:", self.best_metric)
 
+    @timed(stdout=False, return_time=True)
     def train(self):
         self.model.train()
         correct = 0