From b05ee6ad08dd9b4fc32b3fa4358c1e4d4f0d3df3 Mon Sep 17 00:00:00 2001 From: Tom White Date: Mon, 31 Oct 2016 00:12:28 +1300 Subject: [PATCH] Added --save-every-epoch option --- enhance.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/enhance.py b/enhance.py index e1cfe5e..34bcfd0 100644 --- a/enhance.py +++ b/enhance.py @@ -39,6 +39,7 @@ add_arg('--scales', default=2, type=int, help='How ma add_arg('--model', default='medium', type=str, help='Name of the neural network to load/save.') add_arg('--model', default='small', type=str, help='Name of the neural network to load/save.') add_arg('--train', default=False, type=str, help='File pattern to load for training.') +add_arg('--save-every-epoch', default=False, action='store_true', help='Save generator after every training epoch.') add_arg('--batch-resolution', default=192, type=int, help='Resolution of images in training batch.') add_arg('--batch-size', default=15, type=int, help='Number of images per training batch.') add_arg('--buffer-size', default=1500, type=int, help='Total image fragments kept in cache.') @@ -484,6 +485,8 @@ class NeuralEnhancer(object): print(' - adversary mode: generator engaging discriminator.') self.model.adversary_weight.set_value(args.adversary_weight) running = None + if args.save_every_epoch: + self.model.save_generator() except KeyboardInterrupt: pass