Training runs are currently not deterministic

Author: alexkrz

README.md

@@ -72,3 +72,4 @@ Checkpoints for MagFace can be downloaded at: [MagFace Repository](https://githu
 - [x] Add ElasticFace header
 - [x] Remove `mxnet`dependency. Therefore it is necessary to convert the datasets.
 - [ ] Compare MagFace training to official Magface code
+- [ ] Make training runs deterministic at same seed

src/datamodule_hf.py

@@ -62,6 +62,7 @@ def train_dataloader(self):
             batch_size=self.hparams.batch_size,
             num_workers=self.hparams.num_workers,
             shuffle=True,
+            # generator=torch.Generator().manual_seed(42),
         )
 
 

src/pl_module.py

@@ -58,13 +58,16 @@ def training_step(self, batch, batch_idx):
         # logits vector describes the probability for each image to belong to one of n_classes
         loss = self.criterion(logits, targets)
         optimizer_lr = self.optimizers().optimizer.param_groups[0]["lr"]
-        log_dict = {
-            # "step": float(self.current_epoch),  # Overwrite step to plot epochs on x-axis
-            "loss": loss,
-            "optimizer_lr": optimizer_lr,
-            "max_ampl": max_ampl.item(),
-        }
-        self.log_dict(log_dict, on_step=True)
+        self.log("loss", loss, prog_bar=True)
+        self.log("optimizer_lr", optimizer_lr)
+        self.log("max_ampl", max_ampl.item())
+        # log_dict = {
+        #     # "step": float(self.current_epoch),  # Overwrite step to plot epochs on x-axis
+        #     "loss": loss,
+        #     "optimizer_lr": optimizer_lr,
+        #     "max_ampl": max_ampl.item(),
+        # }
+        # self.log_dict(log_dict, on_step=True)
         return loss
 
     def configure_optimizers(self):

train.py

@@ -61,8 +61,11 @@ def main(
     results_dir: str,
     version: Optional[int] = None,
 ):
-    # 1. Set fixed seed
+    # 1. Set fixed seed and flags for deterministic behavior
     pl.seed_everything(cfg.seed)
+    torch.use_deterministic_algorithms(True)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
 
     # 2. Assign datamodule and pl_module
     datamodule = datamodule