def create_model_fn(args):
net = nn.Sequential(
LambdaLayer(lambda x: 2 * (256 / 255) * x - 1),
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 = AutoregressiveSubsetFlow2d(
base_shape=(
3,
32,
32,
),
transforms=[
MultivariateMOLAutoregressiveSubsetTransform2d(
net,
channels=3,
num_mixtures=args.nr_logistic_mix,
num_bins=256,
mean_lambd=lambda x: 2 * (x * 256 / 255) - 1)
])
return model
def create_model_fn(args):
transforms = []
for i in range(args.num_flows):
transforms += [
QuadraticSplineAutoregressiveSubsetTransform2d(
DropPixelCNN(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],
dropout=args.dropout[i],
init_transforms=lambda x: 2 * x - 1),
num_bins=args.num_bins[i])
]
model = AutoregressiveSubsetFlow2d(base_shape=(
3,
32,
32,
),
transforms=transforms)
return model
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
setup.register_data(data.train, data.test)
###################
## 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 ##
#######################
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)
#######################
## Specify optimizer ##
#######################
optimizer = optim.Adam(model.parameters(), lr=args.lr)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
milestones=args.milestones,
gamma=args.gamma)
setup.register_optimizer(optimizer, scheduler)
data = CategoricalCIFAR10()
setup.register_data(data.train, data.test)
###################
## Specify model ##
###################
net = nn.Sequential(LambdaLayer(lambda x: 2 * (256 / 255) * x - 1),
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 = AutoregressiveSubsetFlow2d(base_shape = (3,32,32,),
transforms = [
MultivariateMOLAutoregressiveSubsetTransform2d(net,
channels=3,
num_mixtures=args.nr_logistic_mix,
num_bins=256,
mean_lambd=lambda x: 2*(x*256/255)-1)
])
setup.register_model(model)
#######################
## Specify optimizer ##
#######################
optimizer = optim.Adam(model.parameters(), lr=args.lr)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=args.gamma)
setup.register_optimizer(optimizer, scheduler)