def construct_model_spec(scale_init=2, no_of_layers=8, add_scaling=True):
global layers
num_scales = scale_init
for scale in range(num_scales - 1):
layers.append(
nn.CouplingLayer('checkerboard0',
name='Checkerboard%d_1' % scale,
num_residual_blocks=no_of_layers,
scaling=add_scaling))
layers.append(
nn.CouplingLayer('checkerboard1',
name='Checkerboard%d_2' % scale,
num_residual_blocks=no_of_layers,
scaling=add_scaling))
layers.append(
nn.CouplingLayer('checkerboard0',
name='Checkerboard%d_3' % scale,
num_residual_blocks=no_of_layers,
scaling=add_scaling))
layers.append(nn.SqueezingLayer(name='Squeeze%d' % scale))
layers.append(
nn.CouplingLayer('channel0',
name='Channel%d_1' % scale,
num_residual_blocks=no_of_layers,
scaling=add_scaling))
layers.append(
nn.CouplingLayer('channel1',
name='Channel%d_2' % scale,
num_residual_blocks=no_of_layers,
scaling=add_scaling))
layers.append(
nn.CouplingLayer('channel0',
name='Channel%d_3' % scale,
num_residual_blocks=no_of_layers,
scaling=add_scaling))
layers.append(nn.FactorOutLayer(scale, name='FactorOut%d' % scale))
# # final layer
scale = num_scales - 1
layers.append(
nn.CouplingLayer('checkerboard0',
name='Checkerboard%d_1' % scale,
num_residual_blocks=no_of_layers,
scaling=add_scaling))
layers.append(
nn.CouplingLayer('checkerboard1',
name='Checkerboard%d_2' % scale,
num_residual_blocks=no_of_layers,
scaling=add_scaling))
layers.append(
nn.CouplingLayer('checkerboard0',
name='Checkerboard%d_3' % scale,
num_residual_blocks=no_of_layers,
scaling=add_scaling))
layers.append(
nn.CouplingLayer('checkerboard1',
name='Checkerboard%d_4' % scale,
num_residual_blocks=no_of_layers,
scaling=add_scaling))
layers.append(nn.FactorOutLayer(scale, name='FactorOut%d' % scale))
def construct_model_spec():
global layers
num_scales = 2
for scale in range(num_scales - 1):
layers.append(
nn.CouplingLayer('checkerboard0', name='Checkerboard%d_1' % scale))
layers.append(
nn.CouplingLayer('checkerboard1', name='Checkerboard%d_2' % scale))
layers.append(
nn.CouplingLayer('checkerboard0', name='Checkerboard%d_3' % scale))
#layers.append(nn.CouplingLayer('checkerboard1', name='Checkerboard%d_4' % scale))
#layers.append(nn.CouplingLayer('checkerboard0', name='Checkerboard%d_5' % scale))
#layers.append(nn.CouplingLayer('checkerboard1', name='Checkerboard%d_6' % scale))
layers.append(nn.SqueezingLayer(name='Squeeze%d' % scale))
layers.append(nn.CouplingLayer('channel0', name='Channel%d_1' % scale))
layers.append(nn.CouplingLayer('channel1', name='Channel%d_2' % scale))
layers.append(nn.CouplingLayer('channel0', name='Channel%d_3' % scale))
layers.append(nn.FactorOutLayer(scale, name='FactorOut%d' % scale))
# final layer
scale = num_scales - 1
layers.append(
nn.CouplingLayer('checkerboard0', name='Checkerboard%d_1' % scale))
layers.append(
nn.CouplingLayer('checkerboard1', name='Checkerboard%d_2' % scale))
layers.append(
nn.CouplingLayer('checkerboard0', name='Checkerboard%d_3' % scale))
layers.append(
nn.CouplingLayer('checkerboard1', name='Checkerboard%d_4' % scale))
# layers.append(nn.SqueezingLayer(name='Squeeze%d' % scale))
# layers.append(nn.CouplingLayer('channel0', name='Channel%d_1' % scale))
# layers.append(nn.CouplingLayer('channel1', name='Channel%d_2' % scale))
# layers.append(nn.CouplingLayer('channel0', name='Channel%d_3' % scale))
layers.append(nn.FactorOutLayer(scale, name='FactorOut%d' % scale))