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