fanwuyang
/
neural-enhance-super-resolution
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #18 from alexjc/training

Browse Source 
Training Improvements
main
Alex J. Champandard 9 years ago committed by GitHub
parent a9b0cd9887 17fcad8d28
commit a5ad2c25e6
1 changed files with 44 additions and 38 deletions
Show Stats Download Patch File Download Diff File

72
enhance.py
Unescape View File

@ -38,12 +38,13 @@ add_arg('files', nargs='*', default=[])
add_arg('--scales', default=2, type=int, help='How many times to perform 2x upsampling.') add_arg('--scales', default=2, type=int, help='How many times to perform 2x upsampling.')
add_arg('--model', default='small', 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('--train', default=False, type=str, help='File pattern to load for training.')
add_arg('--batch-resolution', default=192, type=int, help='Resolution of images in training batch.') add_arg('--epochs', default=10, type=int, help='Total number of iterations in training.')
add_arg('--epoch-size', default=72, type=int, help='Number of batches trained in an epoch.')
add_arg('--save-every', default=10, type=int, help='Save generator after every training epoch.')
add_arg('--batch-shape', 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('--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.') add_arg('--buffer-size', default=1500, type=int, help='Total image fragments kept in cache.')
add_arg('--buffer-similar', default=5, type=int, help='Fragments cached for each image loaded.') add_arg('--buffer-similar', default=5, type=int, help='Fragments cached for each image loaded.')
add_arg('--epochs', default=10, type=int, help='Total number of iterations in training.')
add_arg('--epoch-size', default=72, type=int, help='Number of batches trained in an epoch.')
add_arg('--learning-rate', default=1E-4, type=float, help='Parameter for the ADAM optimizer.') add_arg('--learning-rate', default=1E-4, type=float, help='Parameter for the ADAM optimizer.')
add_arg('--learning-period', default=50, type=int, help='How often to decay the learning rate.') add_arg('--learning-period', default=50, type=int, help='How often to decay the learning rate.')
add_arg('--learning-decay', default=0.5, type=float, help='How much to decay the learning rate.') add_arg('--learning-decay', default=0.5, type=float, help='How much to decay the learning rate.')
@ -128,8 +129,10 @@ class DataLoader(threading.Thread):
self.data_ready = threading.Event() self.data_ready = threading.Event()
self.data_copied = threading.Event() self.data_copied = threading.Event()
self.resolution = args.batch_resolution self.orig_shape, self.seed_shape = args.batch_shape, int(args.batch_shape / 2**args.scales)
self.buffer = np.zeros((args.buffer_size, 3, self.resolution, self.resolution), dtype=np.float32)
self.orig_buffer = np.zeros((args.buffer_size, 3, self.orig_shape, self.orig_shape), dtype=np.float32)
self.seed_buffer = np.zeros((args.buffer_size, 3, self.seed_shape, self.seed_shape), dtype=np.float32)
self.files = glob.glob(args.train) self.files = glob.glob(args.train)
if len(self.files) == 0: if len(self.files) == 0:
error("There were no files found to train from searching for `{}`".format(args.train), error("There were no files found to train from searching for `{}`".format(args.train),
@ -157,27 +160,31 @@ class DataLoader(threading.Thread):
for _ in range(args.buffer_similar): for _ in range(args.buffer_similar):
copy = img[:,::-1] if random.choice([True, False]) else img copy = img[:,::-1] if random.choice([True, False]) else img
h = random.randint(0, copy.shape[0] - self.resolution) h = random.randint(0, copy.shape[0] - self.orig_shape)
w = random.randint(0, copy.shape[1] - self.resolution) w = random.randint(0, copy.shape[1] - self.orig_shape)
copy = copy[h:h+self.resolution, w:w+self.resolution] copy = copy[h:h+self.orig_shape, w:w+self.orig_shape]
while len(self.available) == 0: while len(self.available) == 0:
self.data_copied.wait() self.data_copied.wait()
self.data_copied.clear() self.data_copied.clear()
i = self.available.pop() i = self.available.pop()
self.buffer[i] = np.transpose(copy / 255.0 - 0.5, (2, 0, 1)) self.orig_buffer[i] = np.transpose(copy / 255.0 - 0.5, (2, 0, 1))
seed = scipy.misc.imresize(copy, size=(self.seed_shape, self.seed_shape), interp='bilinear')
self.seed_buffer[i] = np.transpose(seed / 255.0 - 0.5, (2, 0, 1))
self.ready.add(i) self.ready.add(i)
if len(self.ready) >= args.batch_size: if len(self.ready) >= args.batch_size:
self.data_ready.set() self.data_ready.set()
def copy(self, output): def copy(self, origs_out, seeds_out):
self.data_ready.wait() self.data_ready.wait()
self.data_ready.clear() self.data_ready.clear()
for i, j in enumerate(random.sample(self.ready, args.batch_size)): for i, j in enumerate(random.sample(self.ready, args.batch_size)):
output[i] = self.buffer[j] origs_out[i] = self.orig_buffer[j]
seeds_out[i] = self.seed_buffer[j]
self.available.add(j) self.available.add(j)
self.data_copied.set() self.data_copied.set()
@ -211,12 +218,8 @@ class Model(object):
def __init__(self): def __init__(self):
self.network = collections.OrderedDict() self.network = collections.OrderedDict()
if args.train:
self.network['img'] = InputLayer((None, 3, None, None)) self.network['img'] = InputLayer((None, 3, None, None))
self.network['seed'] = PoolLayer(self.network['img'], pool_size=2**args.scales, mode='average_exc_pad') self.network['seed'] = InputLayer((None, 3, None, None))
else:
self.network['img'] = InputLayer((None, 3, None, None))
self.network['seed'] = self.network['img']
config, params = self.load_model() config, params = self.load_model()
self.setup_generator(self.last_layer(), config) self.setup_generator(self.last_layer(), config)
@ -378,10 +381,10 @@ class Model(object):
def compile(self): def compile(self):
# Helper function for rendering test images during training, or standalone non-training mode. # Helper function for rendering test images during training, or standalone non-training mode.
input_tensor = T.tensor4() input_tensor, seed_tensor = T.tensor4(), T.tensor4()
input_layers = {self.network['img']: input_tensor} input_layers = {self.network['img']: input_tensor, self.network['seed']: seed_tensor}
output = lasagne.layers.get_output([self.network[k] for k in ['img', 'seed', 'out']], input_layers, deterministic=True) output = lasagne.layers.get_output([self.network[k] for k in ['seed', 'out']], input_layers, deterministic=True)
self.predict = theano.function([input_tensor], output) self.predict = theano.function([seed_tensor], output)
if not args.train: return if not args.train: return
@ -407,7 +410,7 @@ class Model(object):
# Combined Theano function for updating both generator and discriminator at the same time. # Combined Theano function for updating both generator and discriminator at the same time.
updates = collections.OrderedDict(list(gen_updates.items()) + list(disc_updates.items())) updates = collections.OrderedDict(list(gen_updates.items()) + list(disc_updates.items()))
self.fit = theano.function([input_tensor], gen_losses + [disc_out.mean(axis=(1,2,3))], updates=updates) self.fit = theano.function([input_tensor, seed_tensor], gen_losses + [disc_out.mean(axis=(1,2,3))], updates=updates)
@ -448,10 +451,12 @@ class NeuralEnhancer(object):
if t_cur % args.learning_period == 0: l_r *= args.learning_decay if t_cur % args.learning_period == 0: l_r *= args.learning_decay
def train(self): def train(self):
images = np.zeros((args.batch_size, 3, args.batch_resolution, args.batch_resolution), dtype=np.float32) seed_size = int(args.batch_shape / 2**args.scales)
images = np.zeros((args.batch_size, 3, args.batch_shape, args.batch_shape), dtype=np.float32)
seeds = np.zeros((args.batch_size, 3, seed_size, seed_size), dtype=np.float32)
learning_rate = self.decay_learning_rate() learning_rate = self.decay_learning_rate()
try: try:
running, start = None, time.time() average, start = None, time.time()
for epoch in range(args.epochs): for epoch in range(args.epochs):
total, stats = None, None total, stats = None, None
l_r = next(learning_rate) l_r = next(learning_rate)
@ -459,18 +464,18 @@ class NeuralEnhancer(object):
if epoch >= args.discriminator_start: self.model.disc_lr.set_value(l_r) if epoch >= args.discriminator_start: self.model.disc_lr.set_value(l_r)
for _ in range(args.epoch_size): for _ in range(args.epoch_size):
self.thread.copy(images) self.thread.copy(images, seeds)
output = self.model.fit(images) output = self.model.fit(images, seeds)
losses = np.array(output[:3], dtype=np.float32) losses = np.array(output[:3], dtype=np.float32)
stats = (stats + output[3]) if stats is not None else output[3] stats = (stats + output[3]) if stats is not None else output[3]
total = total + losses if total is not None else losses total = total + losses if total is not None else losses
l = np.sum(losses) l = np.sum(losses)
assert not np.isnan(losses).any() assert not np.isnan(losses).any()
running = l if running is None else running * 0.95 + 0.05 * l average = l if average is None else average * 0.95 + 0.05 * l
print('↑' if l > running else '↓', end='', flush=True) print('↑' if l > average else '↓', end='', flush=True)
orign, scald, repro = self.model.predict(images) scald, repro = self.model.predict(seeds)
self.show_progress(orign, scald, repro) self.show_progress(images, scald, repro)
total /= args.epoch_size total /= args.epoch_size
stats /= args.epoch_size stats /= args.epoch_size
totals, labels = [sum(total)] + list(total), ['total', 'prcpt', 'smthn', 'advrs'] totals, labels = [sum(total)] + list(total), ['total', 'prcpt', 'smthn', 'advrs']
@ -481,9 +486,12 @@ class NeuralEnhancer(object):
real, fake = stats[:args.batch_size], stats[args.batch_size:] real, fake = stats[:args.batch_size], stats[args.batch_size:]
print(' - discriminator', real.mean(), len(np.where(real > 0.5)[0]), fake.mean(), len(np.where(fake < -0.5)[0])) print(' - discriminator', real.mean(), len(np.where(real > 0.5)[0]), fake.mean(), len(np.where(fake < -0.5)[0]))
if epoch == args.adversarial_start-1: if epoch == args.adversarial_start-1:
print(' - adversary mode: generator engaging discriminator.') print(' - generator now optimizing against discriminator.')
self.model.adversary_weight.set_value(args.adversary_weight) self.model.adversary_weight.set_value(args.adversary_weight)
running = None running = None
if (epoch+1) % args.save_every == 0:
print(' - saving current generator layers to disk...')
self.model.save_generator()
except KeyboardInterrupt: except KeyboardInterrupt:
pass pass
@ -505,11 +513,9 @@ if __name__ == "__main__":
if args.train: if args.train:
enhancer.train() enhancer.train()
else:
for filename in args.files: for filename in args.files:
print(filename) print(filename)
out = enhancer.process(scipy.ndimage.imread(filename, mode='RGB')) out = enhancer.process(scipy.ndimage.imread(filename, mode='RGB'))
out.save(os.path.splitext(filename)[0]+'_ne%ix.png'%(2**args.scales)) out.save(os.path.splitext(filename)[0]+'_ne%ix.png'%(2**args.scales))
if args.files:
print(ansi.ENDC) print(ansi.ENDC)

Write Preview
Loading…
Cancel
Save