def construct_discrete_model():
chain = []
for i in range(args.depth):
if args.glow: chain.append(layers.BruteForceLayer(2))
chain.append(layers.CouplingLayer(2, swap=i % 2 == 0))
return layers.SequentialFlow(chain)
n_dist=args.n_dist,
n_power_series=args.n_power_series,
exact_trace=args.exact_trace,
brute_force=args.brute_force,
n_samples=args.n_samples,
neumann_grad=False,
grad_in_forward=False,
)
)
if args.actnorm: blocks.append(layers.ActNorm1d(2))
if args.batchnorm: blocks.append(layers.MovingBatchNorm1d(2))
model = layers.SequentialFlow(blocks).to(device)
elif args.arch == 'realnvp':
blocks = []
for _ in range(args.nblocks):
blocks.append(layers.CouplingLayer(2, swap=False))
blocks.append(layers.CouplingLayer(2, swap=True))
if args.actnorm: blocks.append(layers.ActNorm1d(2))
if args.batchnorm: blocks.append(layers.MovingBatchNorm1d(2))
model = layers.SequentialFlow(blocks).to(device)
logger.info(model)
logger.info("Number of trainable parameters: {}".format(count_parameters(model)))
#genGen = Generator(16)
#genGen = Generator(128)
genGen = Generator(8)
#optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
