Merge pull request #110 from yasserben/reproducibility

ensure reproduciblity & consistent restart on Tensorboard side

Author: ClementPinard

main.py

@@ -35,6 +35,12 @@
 group.add_argument('--split-value', default=0.8, type=float,
                    help='test-val split proportion between 0 (only test) and 1 (only train), '
                         'will be overwritten if a split file is set')
+parser.add_argument(
+    "--split-seed",
+    type=int,
+    default=None,
+    help="Seed the train-val split to enforce reproducibility (consistent restart too)",
+)
 parser.add_argument('--arch', '-a', metavar='ARCH', default='flownets',
                     choices=model_names,
                     help='model architecture, overwritten if pretrained is specified: ' +
@@ -79,8 +85,10 @@
                     help='value by which flow will be divided. Original value is 20 but 1 with batchNorm gives good results')
 parser.add_argument('--milestones', default=[100,150,200], metavar='N', nargs='*', help='epochs at which learning rate is divided by 2')
 
+
+
 best_EPE = -1
-n_iter = 0
+n_iter = int(start_epoch)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
@@ -102,6 +110,9 @@ def main():
     if not os.path.exists(save_path):
         os.makedirs(save_path)
 
+    if args.seed_split is not None:
+        np.random.seed(args.seed_split)
+
     train_writer = SummaryWriter(os.path.join(save_path,'train'))
     test_writer = SummaryWriter(os.path.join(save_path,'test'))
     output_writers = []
@@ -294,7 +305,7 @@ def validate(val_loader, model, epoch, output_writers):
         end = time.time()
 
         if i < len(output_writers):  # log first output of first batches
-            if epoch == 0:
+            if epoch == args.start_epoch:
                 mean_values = torch.tensor([0.45,0.432,0.411], dtype=input.dtype).view(3,1,1)
                 output_writers[i].add_image('GroundTruth', flow2rgb(args.div_flow * target[0], max_value=10), 0)
                 output_writers[i].add_image('Inputs', (input[0,:3].cpu() + mean_values).clamp(0,1), 0)