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

Improve console output.

Browse Source
main
Alex J. Champandard 9 years ago
parent 8770113f9a
commit 4d82c7f157
1 changed files with 19 additions and 12 deletions
Show Stats Download Patch File Download Diff File

29
enhance.py
Unescape View File

@ -122,7 +122,7 @@ class DataLoader(threading.Thread):
def run(self): def run(self):
files, cache = glob.glob('train/*.jpg'), {} files, cache = glob.glob('train/*.jpg'), {}
while True: while True:
# random.shuffle(files) random.shuffle(files)
for i, f in enumerate(files[:args.batch_size]): for i, f in enumerate(files[:args.batch_size]):
filename = os.path.join(self.cwd, f) filename = os.path.join(self.cwd, f)
try: try:
@ -133,9 +133,9 @@ class DataLoader(threading.Thread):
files.remove(f) files.remove(f)
continue continue
# if random.choice([True, False]): img[:,:] = img[:,::-1] if random.choice([True, False]): img[:,:] = img[:,::-1]
h = (img.shape[0] - self.resolution) // 2 # random.randint(0, img.shape[0] - self.resolution) h = random.randint(0, img.shape[0] - self.resolution)
w = (img.shape[1] - self.resolution) // 2 # random.randint(0, img.shape[1] - self.resolution) w = random.randint(0, img.shape[1] - self.resolution)
img = img[h:h+self.resolution, w:w+self.resolution] img = img[h:h+self.resolution, w:w+self.resolution]
self.images[i] = np.transpose(img / 255.0 - 0.5, (2, 0, 1)) self.images[i] = np.transpose(img / 255.0 - 0.5, (2, 0, 1))
@ -187,6 +187,8 @@ class Model(object):
self.setup_perceptual(concatenated) self.setup_perceptual(concatenated)
self.load_perceptual() self.load_perceptual()
self.setup_discriminator() self.setup_discriminator()
self.load_generator()
self.compile() self.compile()
#------------------------------------------------------------------------------------------------------------------ #------------------------------------------------------------------------------------------------------------------
@ -285,14 +287,18 @@ class Model(object):
def save_generator(self): def save_generator(self):
def cast(p): return p.get_value().astype(np.float16) def cast(p): return p.get_value().astype(np.float16)
params = {k: [cast(p) for p in l.get_params()] for (k, l) in self.list_generator_layers()} params = {k: [cast(p) for p in l.get_params()] for (k, l) in self.list_generator_layers()}
pickle.dump(params, bz2.open(args.model % 2**args.scales, 'wb')) filename = args.model % 2**args.scales
pickle.dump(params, bz2.open(filename, 'wb'))
print(' - Saved model as `{}` after training.'.format(filename))
def load_generator(self): def load_generator(self):
filename = args.model % 2**args.scales filename = args.model % 2**args.scales
if not os.path.exists(filename): return if not os.path.exists(filename): return
params = pickle.load(bz2.open(filename, 'rb')) params = pickle.load(bz2.open(filename, 'rb'))
for k, l in self.list_generator_layers(): for k, l in self.list_generator_layers():
if k not in params: continue
(p.set_value(v) for p, v in zip(l.get_params(), params[k])) (p.set_value(v) for p, v in zip(l.get_params(), params[k]))
print(' - Loaded file `{}` with trained model.'.format(filename))
#------------------------------------------------------------------------------------------------------------------ #------------------------------------------------------------------------------------------------------------------
# Training & Loss Functions # Training & Loss Functions
@ -367,15 +373,13 @@ class NeuralEnhancer(object):
self.imsave('valid/%03i_repro.png' % i, repro[i]) self.imsave('valid/%03i_repro.png' % i, repro[i])
def train(self): def train(self):
self.model.load_generator()
images = np.zeros((args.batch_size, 3, args.batch_resolution, args.batch_resolution), dtype=np.float32) images = np.zeros((args.batch_size, 3, args.batch_resolution, args.batch_resolution), dtype=np.float32)
l_min, l_max, l_mult = 1E-7, 1E-3, 0.2 l_min, l_max, l_mult = 1E-7, 1E-3, 0.2
t_cur, t_i, t_mult = 120, 150, 1 t_cur, t_i, t_mult = 120, 150, 1
i, running = 0, None i, running, start = 0, None, time.time()
for k in range(args.epochs): for k in range(args.epochs):
total, start = None, time.time() total = None
for _ in range(args.epoch_size): for _ in range(args.epoch_size):
i += 1 i += 1
l_r = l_min + 0.5 * (l_max - l_min) * (1.0 + math.cos(t_cur / t_i * math.pi)) l_r = l_min + 0.5 * (l_max - l_min) * (1.0 + math.cos(t_cur / t_i * math.pi))
@ -400,14 +404,17 @@ class NeuralEnhancer(object):
total /= args.epoch_size total /= args.epoch_size
totals, labels = [sum(total)] + list(total), ['total', 'prcpt', 'smthn', 'advrs'] totals, labels = [sum(total)] + list(total), ['total', 'prcpt', 'smthn', 'advrs']
losses = ['{}{}{}={:4.2e}'.format(ansi.WHITE_B, k, ansi.ENDC, v) for k, v in zip(labels, totals)] losses = ['{}{}{}={:4.2e}'.format(ansi.WHITE_B, k, ansi.ENDC, v) for k, v in zip(labels, totals)]
print('\rEpoch #{} in {:4.1f}s{}'.format(k+1, time.time()-start, ' '*args.epoch_size)) print('\rEpoch #{} at {:4.1f}s{}'.format(k+1, time.time()-start, ' '*args.epoch_size))
print(' - losses {}\n'.format(' '.join(losses))) print(' - losses {}'.format(' '.join(losses)))
# print(stats[:args.batch_size].mean(), stats[args.batch_size:].mean()) # print(stats[:args.batch_size].mean(), stats[args.batch_size:].mean())
# if k == 0: self.model.disc_lr.set_value(l_r) # if k == 0: self.model.disc_lr.set_value(l_r)
# if k == 1: self.model.adversary_weight.set_value(args.adversary_weight) # if k == 1: self.model.adversary_weight.set_value(args.adversary_weight)
print('\n{}Trained {}x super-resolution for {} epochs.{}'\
.format(ansi.CYAN_B, 2**args.scales, args.epochs, ansi.CYAN))
self.model.save_generator() self.model.save_generator()
print(ansi.ENDC)
if __name__ == "__main__": if __name__ == "__main__":

Write Preview
Loading…
Cancel
Save