fix 2d augment

Author: AndrewBoessen

train.py

@@ -173,7 +173,8 @@ def train_epoch(cfg, model, train_data_loader, loss, optimzer):
 
         with torch.no_grad():
             batch_gt[:,:,:,2] = batch_gt[:,:,:,2] - batch_gt[:,0:1,0:1,2] # Place the depth of first frame root to 0
-            batch_input = cfg.aug.augment2D(batch_input, mask=cfg.mask, noise=cfg.noise)
+            if cfg.noise or cfg.mask:
+                batch_input = cfg.aug.augment2D(batch_input, mask=cfg.mask, noise=cfg.noise)
 
         predicted_3d_pos = model(batch_input) # (N,T,17,3)
 
@@ -287,7 +288,7 @@ def train(args, cfg):
 
         # Load checkpoint if given
         if checkpoint:
-            st = checkpoint['epoch']
+            st = checkpoint['epoch'] if not cfg.finetune else 0
             if 'optimizer' in checkpoint and checkpoint['optimizer'] != None:
                 optimizer.load_state_dict(checkpoint['optimizer'])
             else:
@@ -299,7 +300,7 @@ def train(args, cfg):
         cfg.mask = (cfg.mask_ratio > 0 and cfg.mask_T_ratio > 0)
         if cfg.mask or cfg.noise:
             cfg.aug = Augmenter2D(cfg) # Data Augmentation: flip and add noise
-        
+
         for epoch in range(st, args.epochs or cfg.epochs): # Start training
             print("Training Epoch %d" % (epoch + 1))
             start_time = time()