diff --git a/MARBLE/main.py b/MARBLE/main.py index bdae632d..9469f16e 100644 --- a/MARBLE/main.py +++ b/MARBLE/main.py @@ -348,7 +348,9 @@ def fit(self, data, outdir=None, verbose=False): print("\n---- Training network ...") - time = datetime.now().strftime("%Y%m%d-%H%M%S") + self.timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + + print("\n---- Timestamp: {}".format(self.timestamp)) # load to gpu (if possible) # pylint: disable=self-cls-assignment @@ -385,7 +387,7 @@ def fit(self, data, outdir=None, verbose=False): if best_loss == -1 or (val_loss < best_loss): outdir = self.save_model( - optimizer, self.losses, outdir=outdir, best=True, timestamp=time + optimizer, self.losses, outdir=outdir, best=True, timestamp=self.timestamp ) best_loss = val_loss print(" *", end="") @@ -398,8 +400,8 @@ def fit(self, data, outdir=None, verbose=False): self.losses["test_loss"].append(test_loss) - self.save_model(optimizer, self.losses, outdir=outdir, best=False, timestamp=time) - self.load_model(os.path.join(outdir, f"best_model_{time}.pth")) + self.save_model(optimizer, self.losses, outdir=outdir, best=False, timestamp=self.timestamp) + self.load_model(os.path.join(outdir, f"best_model_{self.timestamp}.pth")) def load_model(self, loadpath): """Load model.