def create_model_fn(args): model = AutoregressiveSubsetFlow2d(base_shape = (3,32,32,), transforms = [ LinearSplineAutoregressiveSubsetTransform2d(PixelCNN(3, num_params=args.num_bins, num_blocks=args.num_blocks, filters=args.filters, kernel_size=args.kernel_size, kernel_size_in=args.kernel_size_in, output_filters=args.output_filters, init_transforms=lambda x: 2*x-1), num_bins=args.num_bins), ]) return model
def create_model_fn(args): if args.bin_cond: lambd = lambda x: 2 * quantize(x, 256) - 1 else: lambd = lambda x: 2 * x - 1 net = PixelCNN(3, num_params=2 * args.num_bins + 1, num_blocks=args.num_blocks, filters=args.filters, kernel_size=args.kernel_size, kernel_size_in=args.kernel_size_in, output_filters=args.output_filters, init_transforms=lambd) model = UniformlyDequantizedFlow( base_dist=StandardUniform(( 3, 32, 32, )), transforms=QuadraticSplineAutoregressiveTransform2d( net, num_bins=args.num_bins)) return model
################### ## Specify model ## ################### model = AutoregressiveSubsetFlow2d( base_shape=( 3, 32, 32, ), transforms=[ QuadraticSplineAutoregressiveSubsetTransform2d(PixelCNN( 3, num_params=2 * args.num_bins + 1, num_blocks=args.num_blocks, filters=args.filters, kernel_size=args.kernel_size, kernel_size_in=args.kernel_size_in, output_filters=args.output_filters, init_transforms=lambda x: 2 * x - 1), num_bins=args.num_bins), ]) setup.register_model(model) ####################### ## Specify optimizer ## ####################### optimizer = optim.Adam(model.parameters(), lr=args.lr) scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
setup.register_data(data.train, data.test) ################### ## Specify model ## ################### transforms = [] for i in range(args.num_flows): transforms += [ QuadraticSplineAutoregressiveSubsetTransform2d( PixelCNN(3, num_params=2 * args.num_bins[i] + 1, num_blocks=args.num_blocks[i], filters=args.filters[i], kernel_size=args.kernel_size[i], kernel_size_in=args.kernel_size_in[i], output_filters=args.output_filters[i], init_transforms=lambda x: 2 * x - 1), num_bins=args.num_bins[i]) ] model = AutoregressiveSubsetFlow2d(base_shape=( 3, 64, 64, ), transforms=transforms) setup.register_model(model)
setup.register_data(data.train, data.test) ################### ## Specify model ## ################### transforms = [] for i in range(args.num_flows): transforms += [ LinearSplineAutoregressiveSubsetTransform2d(PixelCNN( 3, num_params=args.num_bins[i], num_blocks=args.num_blocks[i], filters=args.filters[i], kernel_size=args.kernel_size[i], kernel_size_in=args.kernel_size_in[i], output_filters=args.output_filters[i], init_transforms=lambda x: 2 * x - 1), num_bins=args.num_bins[i]) ] model = AutoregressiveSubsetFlow2d(base_shape=( 3, 64, 64, ), transforms=transforms) setup.register_model(model)
data = CategoricalCIFAR10() setup.register_data(data.train, data.test) ################### ## Specify model ## ################### if args.bin_cond: lambd = lambda x: 2 * quantize(x, 256) - 1 else: lambd = lambda x: 2 * x - 1 net = PixelCNN(3, num_params=2 * args.num_bins + 1, num_blocks=args.num_blocks, filters=args.filters, kernel_size=args.kernel_size, kernel_size_in=args.kernel_size_in, output_filters=args.output_filters, init_transforms=lambd) model = UniformlyDequantizedFlow( base_dist=StandardUniform(( 3, 32, 32, )), transforms=QuadraticSplineAutoregressiveTransform2d( net, num_bins=args.num_bins)) setup.register_model(model)