Skip to content

Commit

Permalink
Add support for PRELU (default), optional load/save.
Browse files Browse the repository at this point in the history
  • Loading branch information
alexjc committed Oct 4, 2016
1 parent 4d82c7f commit b52530c
Showing 1 changed file with 13 additions and 6 deletions.
19 changes: 13 additions & 6 deletions enhance.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
parser = argparse.ArgumentParser(description='Generate a new image by applying style onto a content image.',
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
add_arg = parser.add_argument
add_arg('--load', default=None, action='store_true')
add_arg('--save', default=None, action='store_true')
add_arg('--model', default='ne%ix.pkl.bz2', type=str)
add_arg('--batch-size', default=15, type=int)
add_arg('--batch-resolution', default=256, type=int)
Expand All @@ -42,7 +44,7 @@
add_arg('--generator-residual', default=2, type=int)
add_arg('--perceptual-layer', default='conv2_2', type=str)
add_arg('--perceptual-weight', default=1e0, type=float)
add_arg('--smoothness-weight', default=1e4, type=float)
add_arg('--smoothness-weight', default=1e6, type=float)
add_arg('--adversary-weight', default=0.0, type=float)
add_arg('--scales', default=1, type=int, help='')
add_arg('--device', default='gpu0', type=str, help='Name of the CPU/GPU number to use, for Theano.')
Expand Down Expand Up @@ -198,14 +200,17 @@ def __init__(self):
def last_layer(self):
return list(self.network.values())[-1]

def make_layer(self, input, units, filter_size=(3,3), stride=(1,1), pad=(1,1)):
return ConvLayer(input, units, filter_size=filter_size, stride=stride, pad=pad,
nonlinearity=lasagne.nonlinearities.rectify)
def make_layer(self, input, units, filter_size=(3,3), stride=(1,1), pad=(1,1), nl='prelu'):
conv = ConvLayer(input, units, filter_size=filter_size, stride=stride, pad=pad,
nonlinearity=lasagne.nonlinearities.linear)
if nl == 'relu': conv.nonlinearity = lasagne.nonlinearities.rectify
if nl == 'prelu': conv = lasagne.layers.prelu(conv)
return conv

def make_block(self, name, input, units):
self.network[name+'|Ac'] = ConvLayer(input, units, filter_size=(3,3), stride=(1,1), pad=1)
self.network[name+'|Ac'] = self.make_layer(input, units)
self.network[name+'|An'] = batch_norm(self.last_layer()).input_layer
self.network[name+'|Bc'] = ConvLayer(self.last_layer(), units, filter_size=(3,3), stride=(1,1), pad=1)
self.network[name+'|Bc'] = self.make_layer(self.last_layer(), units)
self.network[name+'|Bn'] = batch_norm(self.last_layer()).input_layer
return ElemwiseSumLayer([input, self.last_layer()]) if args.generator_residual else self.last_layer()

Expand Down Expand Up @@ -285,13 +290,15 @@ def list_generator_layers(self):
yield (name, l)

def save_generator(self):
if not args.save: return
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()}
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):
if not args.load: return
filename = args.model % 2**args.scales
if not os.path.exists(filename): return
params = pickle.load(bz2.open(filename, 'rb'))
Expand Down

0 comments on commit b52530c

Please sign in to comment.