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))