uupdate

requirements.txt

@@ -8,5 +8,5 @@ regex
 overrides
 omegaconf>=2.1
 pyarrow
-colorlog
+colorlog==5.0.1
 hydra-core

torchfly/training/callbacks/evaluation.py

@@ -165,6 +165,7 @@ def record_validation_metrics(self, trainer: Trainer):
     @handle_event(Events.TEST_BEGIN)
     def info_test_begin(self, trainer: Trainer):
         logger.info(f"Test starts! ")
+        os.makedirs("evaluation", exist_ok=True)
         self.eval_start_time = time.time()
 
     @handle_event(Events.TEST_END)

torchfly/utilities/save_config.py

@@ -0,0 +1,5 @@
+from omegaconf import OmegaConf
+
+def save_config(config, filename):
+    with open(filename, "w") as f:
+        OmegaConf.save(config, f)    

tutorials/MNIST/config/training/base.yaml

@@ -8,5 +8,5 @@ training:
     evaluation:
         save_top_k_models: 3
     total_num:
-        epochs: 20
+        epochs: 5
         update_steps: -1 # disabled when total_num.epochs < 0

tutorials/MNIST/main.py

@@ -6,7 +6,9 @@
 from torchfly.flylogger import FlyLogger
 from torchfly.flyconfig import FlyConfig
 from torchfly.training import Trainer
-from torchfly.utils import distributed
+import torchfly.distributed as distributed
+from torchfly.utilities import set_random_seed
+from omegaconf import OmegaConf
 
 from model import CNNFlyModel
 
@@ -23,6 +25,7 @@ def train_loader_fn(self):
         }
 
         with distributed.mutex() as rank:
+
             dataset = datasets.MNIST(os.path.join(self.config.task.datadir, 'MNIST'),
                                      train=True,
                                      download=True,
@@ -57,6 +60,8 @@ def main():
         torch.distributed.init_process_group(backend='nccl', init_method='env://')
 
     config = FlyConfig.load("config/config.yaml")
+    set_random_seed(config.training.random_seed)
+
     data_helper = DataLoaderHelper(config)
     train_dataloader = data_helper.train_loader_fn()
     valid_dataloader = data_helper.valid_loader_fn()
@@ -65,9 +70,14 @@ def main():
 
     trainer = Trainer(config.training, model)
 
-    with FlyLogger(config.flylogger, overwrite=True) as flylogger:
+    model.configure_metrics()
+
+    with FlyLogger(config.flylogger) as flylogger:
+        with open("config.yaml", "w") as f:
+            OmegaConf.save(config, f)    
         trainer.train(train_dataloader, valid_dataloader)
 
 
+
 if __name__ == "__main__":
     main()