def cglow_SR_layer(layer_index,
cond_channels,
upfactor,
parameterize: Parameterize,
depth=4,
nn_ctor=coupling_nn(),
split_axis=-1,
act_norm=True,
name='glow_layer'):
squeeze = Squeeze(name=f'{name}_squeeze')
steps = Flow.uniform(
depth, lambda i: cglow_SR_step(i,
layer_index,
cond_channels,
upfactor,
nn_ctor=nn_ctor,
name=f'{name}_step{i}'))
norm = ActNorm(name=f'{name}_trs_act_norm')
inv_conv = InvertibleConv(name=f'{name}_trs_inv_conv')
layer_steps = [squeeze, norm, inv_conv, steps]
if split_axis is not None:
layer_steps.append(
CondSplit(parameterize,
cond_channels=cond_channels,
split_axis=split_axis,
name=f'{name}_split'))
return Flow(layer_steps)
def cglow_SR_step(depth_index,
layer_index,
cond_channels,
upfactor,
nn_ctor=coupling_nn(),
name='cglow_SR_step'):
norm = ActNorm(name=f'{name}_act_norm')
injector = AffineInjector(layer_index,
cond_channels,
nn_ctor=nn_ctor,
name=f'{name}_affine_injector')
if True:
inv_conv = InvertibleConv(name=f'{name}_inv_conv')
coupling = CondAffineCoupling(layer_index,
cond_channels,
nn_ctor=nn_ctor,
name=f'{name}_cond_affine_coupling',
reverse=False)
flow_steps = [norm, inv_conv, injector, coupling]
else:
coupling = CondAffineCoupling(layer_index,
cond_channels,
nn_ctor=nn_ctor,
name=f'{name}_cond_affine_coupling',
reverse=True)
flow_steps = [norm, injector, coupling]
return Flow(flow_steps)
def glow_step(layer_index, nn_ctor=coupling_nn(), name='glow_step'):
norm = ActNorm(name=f'{name}_act_norm'
) # if act_norm else BatchNorm(name=f'{name}_batch_norm')
invertible_conv = InvertibleConv(name=f'{name}_inv_conv')
affine_coupling = AffineCoupling(layer_index,
nn_ctor=nn_ctor,
name=f'{name}_affine_coupling')
flow_steps = [norm, invertible_conv, affine_coupling]
return Flow(flow_steps)
def glow_layer(layer_index,
parameterize: Parameterize,
depth=4,
nn_ctor=coupling_nn(),
split_axis=-1,
act_norm=True,
name='glow_layer'):
squeeze = Squeeze(name=f'{name}_squeeze')
steps = Flow.uniform(
depth, lambda i: glow_step(
layer_index, nn_ctor=nn_ctor, name=f'{name}_step{i}'))
layer_steps = [squeeze, steps]
if split_axis is not None:
layer_steps.append(
Split(parameterize, split_axis=split_axis, name=f'{name}_split'))
return Flow(layer_steps)
def __init__(self,
dim=2,
input_channels=None,
num_layers=1,
depth=4,
cond_shape=None,
parameterize_ctor=Gaussianize,
coupling_nn_ctor=coupling_nn(),
act_norm=True,
name='glow_flow',
*args,
**kwargs):
"""
Creates a new Glow normalizing flow with the given configuration.
dim : spatial dimension of the input
input_channels : channel dimension of the input; can be provided here
or at a later time to 'initialize'
num_layers : number of "layers" in the multi-scale Glow architecture
depth : number of glow steps per layer
parameterize_ctor : a function () -> Paramterize (see consructor docs for Split)
coupling_nn_ctor : function that constructs a Keras model for affine coupling steps
act_norm : if true, use act norm in Glow layers; otherwise, use batch norm
"""
def _layer(i):
"""Builds layer i; omits split op for final layer"""
assert i < num_layers, f'expected i < {num_layers}; got {i}'
return glow_layer(i,
parameterize_ctor(name=f'{name}_layer{i}_param'),
depth=depth,
nn_ctor=coupling_nn_ctor,
act_norm=act_norm,
split_axis=None if i == num_layers - 1 else -1,
name=f'{name}_layer{i}')
super().__init__(*args, name=name, **kwargs)
self.dim = dim
self.num_layers = num_layers
self.depth = depth
self.cond_shape = cond_shape
self.layers = [_layer(i) for i in range(num_layers)]
self.parameterize = parameterize_ctor()
if input_channels is not None:
self.input_channels = input_channels
input_shape = tf.TensorShape(
(None, *[None for _ in range(self.dim)], self.input_channels))
self.initialize(input_shape)
def __init__(self,
layer,
cond_shape,
input_shape=None,
nn_ctor=coupling_nn(),
name='affine_injector',
*args,
**kwargs):
self.cond_shape = cond_shape
self.layer = layer
self.nn_ctor = nn_ctor
self.nn = None
super().__init__(*args,
input_shape=input_shape,
requires_init=True,
name=name,
**kwargs)
def __init__(self,
dim=2,
input_channels=None,
upfactor=1,
num_layers=1,
depth=4,
parameterize_ctor=Gaussianize(),
coupling_ctor=coupling_nn(),
cond_ctor=cond_nn(),
act_norm=True,
name='cglow_flow',
*args,
**kwargs):
"""
Creates a new Glow normalizing flow with the given configuration.
dim : the spatial dimension of the input x
input_channels : the channel dimension of the input x;
can be provided here or at a later time to 'initialize'
num_layers : number of "layers" in the multi-scale Glow architecture
depth : number of glow steps per layer
parameterize_ctor : a function () -> Paramterize (see consructor docs for Split)
coupling_ctor : function that constructs a Keras model for the coupling networks
cond_ctor : function that constructs a Keras model for conditioning network
act_norm : if true, use act norm in Glow layers; otherwise, use batch norm
"""
self.dim = dim
self.cond_fn = cond_ctor()
self.cond_channels = self.cond_fn.outputs[0].shape[-1]
def _layer(i):
"""Builds layer i; omits split op for final layer"""
assert i < num_layers, f'expected i < {num_layers}; got {i}'
return cglow_SR_layer(
i,
self.cond_channels,
upfactor,
parameterize_ctor(i=i,
name=f'{name}_layer{i}_param',
cond_channels=self.cond_channels),
depth=depth,
nn_ctor=coupling_ctor,
act_norm=act_norm,
split_axis=None if i == num_layers - 1 else -1,
name=f'{name}_layer{i}')
Transform.__init__(self, *args, name=name, **kwargs)
self.num_layers = num_layers
self.depth = depth
self.upfactor = upfactor
self.parameterize = parameterize_ctor(i=num_layers - 1,
cond_channels=self.cond_channels)
self.layers = [_layer(i) for i in range(num_layers)]
if input_channels is not None:
self.input_channels = input_channels
input_shape = tf.TensorShape(
(None, *[None for i in range(self.dim)], self.input_channels))
self.initialize(input_shape)
def get_model(config, restart=False, ckpt_num=0):
dim = config.getint('flow', 'dim')
upfactor = config.getint('flow', 'upfactor')
num_layers = config.getint('flow', 'num_layers')
rundir = get_rundir(config)
inpt_channels = get_inpt_channels(config)
cond_channels = get_cond_channels(config)
kwargs_nn = get_kwargs(config,
keys=[
'dim', 'min_filters', 'max_filters',
'num_blocks', 'kernel_size'
])
kwargs_flow = get_kwargs(config, keys=['upfactor', 'num_layers', 'depth'])
kwargs_cond = get_kwargs(config,
keys=[
'dim', 'cond_filters', 'cond_resblocks',
'cond_blocks', 'kernel_size'
])
kwargs_cond['cond_channels'] = cond_channels
print("Building model:")
coupling_ctor = coupling_nn(**kwargs_nn)
cond_ctor = cond_nn(upfactor=upfactor,
num_layers=num_layers,
**kwargs_cond)
parametrizer = CondGaussianize(**kwargs_nn)
#parametrizer = Gaussianize(**kwargs_nn)
glow = Invert(
CGlowFlowSR(**kwargs_flow,
**kwargs_cond,
coupling_ctor=coupling_ctor,
cond_ctor=cond_ctor,
parameterize_ctor=parametrizer))
learning_rate = float(config.getfloat('training', 'learning_rate'))
num_bins = config.getint('training', 'num_bins')
prior = NormalPrior(loc=0.0, scale=1.0)
model = FlowLVM(glow,
prior,
dim=dim,
num_bins=num_bins,
input_channels=inpt_channels,
cond_channels=cond_channels,
learning_rate=learning_rate,
rundir=rundir)
print("Model built with",
model.param_count().numpy(),
"parameters.",
flush=True)
model._init_checkpoint()
if config.getint('training', 'restart') > 0 or restart:
rundir = get_rundir(config)
if ckpt_num > 0:
ckpt_path = rundir + "model-{}".format(ckpt_num)
else:
print("Trying to restart from latest checkpoint in:", rundir)
ckpt_path = tf.train.latest_checkpoint(rundir)
print("Checkpoint found:", ckpt_path)
model.checkpoint.restore(ckpt_path)
return model