def create_model_fn(args):
if args.bin_cond: lambd = lambda x: 2*quantize(x, 256) - 1
else: lambd = lambda x: 2*(256/255)*x-1
net = nn.Sequential(LambdaLayer(lambd),
PixelCNNpp(nr_resnet=args.nr_resnet, nr_filters=args.nr_filters, nr_scales=args.nr_scales,
input_channels=3, nr_logistic_mix=args.nr_logistic_mix, dropout=args.dropout))
model = UniformlyDequantizedFlow(base_dist = StandardUniform((3,32,32,)),
transforms = MultivariateCMOLAutoregressiveTransform2d(net,
channels=3,
num_mixtures=args.nr_logistic_mix,
num_bins=256,
mean_lambd=lambd))
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
args = parser.parse_args()
setup.prepare_env(args)
##################
## Specify data ##
##################
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,