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 __init__(self, in_channels, num_params, filters=128, num_blocks=15, output_filters=1024, kernel_size=3, kernel_size_in=7, init_transforms=lambda x: 2*x-1):
layers = [LambdaLayer(init_transforms)] +\
[MaskedConv2d(in_channels, 2 * filters, kernel_size=kernel_size_in, padding=kernel_size_in//2, mask_type='A', data_channels=in_channels)] +\
[MaskedResidualBlock2d(filters, data_channels=in_channels, kernel_size=kernel_size) for _ in range(num_blocks)] +\
[nn.ReLU(True), MaskedConv2d(2 * filters, output_filters, kernel_size=1, mask_type='B', data_channels=in_channels)] +\
[nn.ReLU(True), MaskedConv2d(output_filters, num_params * in_channels, kernel_size=1, mask_type='B', data_channels=in_channels)] +\
[ElementwiseParams2d(num_params)]
super(PixelCNN, self).__init__(*layers)
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
## 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 * (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,
setup.prepare_env(args)
##################
## Specify data ##
##################
data = CategoricalCIFAR10()
setup.register_data(data.train, data.test)
###################
## Specify model ##
###################
pixelcnn_pp = 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))
pixelcnn = DropPixelCNN(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,
dropout=args.pixelcnn_dropout,
args = parser.parse_args()
setup.prepare_env(args)
##################
## Specify data ##
##################
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)
#######################