def create_base_transform(i):
if args.base_transform_type == 'rq':
return transforms.PiecewiseRationalQuadraticCouplingTransform(
mask=utils.create_alternating_binary_mask(features=dim,
even=(i % 2 == 0)),
transform_net_create_fn=lambda in_features, out_features: nn_.
ResidualNet(in_features=in_features,
out_features=out_features,
hidden_features=args.hidden_features,
num_blocks=args.num_transform_blocks,
dropout_probability=args.dropout_probability,
use_batch_norm=args.use_batch_norm),
num_bins=args.num_bins,
apply_unconditional_transform=False,
)
elif args.base_transform_type == 'affine':
return transforms.AffineCouplingTransform(
mask=utils.create_alternating_binary_mask(features=dim,
even=(i % 2 == 0)),
transform_net_create_fn=lambda in_features, out_features: nn_.
ResidualNet(in_features=in_features,
out_features=out_features,
hidden_features=args.hidden_features,
num_blocks=args.num_transform_blocks,
dropout_probability=args.dropout_probability,
use_batch_norm=args.use_batch_norm))
else:
raise ValueError
def create_base_transform(i):
if args.base_transform_type == 'affine':
return transforms.AffineCouplingTransform(
mask=utils.create_alternating_binary_mask(features=dim,
even=(i % 2 == 0)),
transform_net_create_fn=lambda in_features, out_features: nn_.
ResidualNet(in_features=in_features,
out_features=out_features,
hidden_features=32,
num_blocks=2,
use_batch_norm=True))
else:
return transforms.PiecewiseRationalQuadraticCouplingTransform(
mask=utils.create_alternating_binary_mask(features=dim,
even=(i % 2 == 0)),
transform_net_create_fn=lambda in_features, out_features: nn_.
ResidualNet(in_features=in_features,
out_features=out_features,
hidden_features=32,
num_blocks=2,
use_batch_norm=True),
tails='linear',
tail_bound=5,
num_bins=args.num_bins,
apply_unconditional_transform=False)
def create_base_transform(i, context_features=None):
if args.prior_type == 'affine-coupling':
return transforms.AffineCouplingTransform(
mask=utils.create_alternating_binary_mask(
features=args.latent_features, even=(i % 2 == 0)),
transform_net_create_fn=lambda in_features, out_features: nn_.
ResidualNet(in_features=in_features,
out_features=out_features,
hidden_features=args.hidden_features,
context_features=context_features,
num_blocks=args.num_transform_blocks,
activation=F.relu,
dropout_probability=args.dropout_probability,
use_batch_norm=args.use_batch_norm))
elif args.prior_type == 'rq-coupling':
return transforms.PiecewiseRationalQuadraticCouplingTransform(
mask=utils.create_alternating_binary_mask(
features=args.latent_features, even=(i % 2 == 0)),
transform_net_create_fn=lambda in_features, out_features: nn_.
ResidualNet(in_features=in_features,
out_features=out_features,
hidden_features=args.hidden_features,
context_features=context_features,
num_blocks=args.num_transform_blocks,
activation=F.relu,
dropout_probability=args.dropout_probability,
use_batch_norm=args.use_batch_norm),
num_bins=args.num_bins,
tails='linear',
tail_bound=args.tail_bound,
apply_unconditional_transform=args.
apply_unconditional_transform,
)
elif args.prior_type == 'affine-autoregressive':
return transforms.MaskedAffineAutoregressiveTransform(
features=args.latent_features,
hidden_features=args.hidden_features,
context_features=context_features,
num_blocks=args.num_transform_blocks,
use_residual_blocks=True,
random_mask=False,
activation=F.relu,
dropout_probability=args.dropout_probability,
use_batch_norm=args.use_batch_norm)
elif args.prior_type == 'rq-autoregressive':
return transforms.MaskedPiecewiseRationalQuadraticAutoregressiveTransform(
features=args.latent_features,
hidden_features=args.hidden_features,
context_features=context_features,
num_bins=args.num_bins,
tails='linear',
tail_bound=args.tail_bound,
num_blocks=args.num_transform_blocks,
use_residual_blocks=True,
random_mask=False,
activation=F.relu,
dropout_probability=args.dropout_probability,
use_batch_norm=args.use_batch_norm)
else:
raise ValueError
def __init__(self, num_input_channels, num_hidden_layers,
num_hidden_channels, num_bins, tail_bound, activation,
dropout_probability, reverse_mask):
def transform_net_create_fn(in_features, out_features):
return ResidualNet(in_features=in_features,
out_features=out_features,
context_features=None,
hidden_features=num_hidden_channels,
num_blocks=num_hidden_layers,
activation=activation(),
dropout_probability=dropout_probability,
use_batch_norm=False)
super().__init__(
num_input_channels=num_input_channels,
flow=PiecewiseRationalQuadraticCouplingTransform(
mask=create_alternating_binary_mask(num_input_channels,
even=reverse_mask),
transform_net_create_fn=transform_net_create_fn,
num_bins=num_bins,
tails='linear',
tail_bound=tail_bound,
# Setting True corresponds to equations (4), (5), (6) in the NSF paper:
apply_unconditional_transform=True))
def create_base_transform(i):
base_transform = transforms.PiecewiseCircularRationalQuadraticCouplingTransform_edge(
mask=utils.create_alternating_binary_mask(feature_dim,
even=(i % 2 == 0)),
transform_net_create_fn=lambda in_features, out_features: nn_.
ResidualNet(in_features=in_features,
out_features=out_features,
hidden_features=args.hidden_features,
context_features=context_dim,
num_blocks=args.num_transform_blocks,
activation=F.relu,
dropout_probability=args.dropout_probability,
use_batch_norm=args.use_batch_norm),
num_bins=args.num_bins,
tails=None,
apply_unconditional_transform=args.apply_unconditional_transform)
return base_transform
# transform_net_create_fn=lambda in_features, out_features: nn_.ResidualNet(
# in_features=in_features,
# out_features=out_features,
# hidden_features=32,
# num_blocks=2,
# use_batch_norm=True
# ),
# num_bins=args.num_bins,
# apply_unconditional_transform=False
# )
# ])
transform = transforms.CompositeTransform([
# transforms.Sigmoid(),
transforms.PiecewiseRationalLinearCouplingTransform(
mask=utils.create_alternating_binary_mask(features=dim, even=True),
transform_net_create_fn=lambda in_features, out_features: nn_.ResidualNet(
in_features=in_features,
out_features=out_features,
hidden_features=32,
num_blocks=2,
use_batch_norm=True
),
tails='linear',
tail_bound=5,
num_bins=args.num_bins,
apply_unconditional_transform=False
),
transforms.PiecewiseRationalLinearCouplingTransform(
mask=utils.create_alternating_binary_mask(features=dim, even=False),
transform_net_create_fn=lambda in_features, out_features: nn_.ResidualNet(
# Generate test grid data
num_points_per_axis = 512
bounds = np.array([[1e-3, 1 - 1e-3], [1e-3, 1 - 1e-3]])
grid_dataset = data_.TestGridDataset(num_points_per_axis=num_points_per_axis,
bounds=bounds)
grid_loader = data.DataLoader(dataset=grid_dataset,
batch_size=1000,
drop_last=False)
# create model
distribution = uniform.TweakedUniform(low=torch.zeros(dim),
high=torch.ones(dim))
transform = transforms.CompositeTransform([
transforms.PiecewiseRationalLinearCouplingTransform(
mask=utils.create_alternating_binary_mask(features=dim,
even=(i % 2 == 0)),
transform_net_create_fn=lambda in_features, out_features: nn_.
ResidualNet(in_features=in_features,
out_features=out_features,
hidden_features=args.hidden_features,
num_blocks=args.num_transform_blocks,
dropout_probability=args.dropout_probability,
use_batch_norm=args.use_batch_norm),
num_bins=args.num_bins,
tails=None,
tail_bound=1,
# apply_unconditional_transform=args.apply_unconditional_transform,
min_bin_width=args.min_bin_width) for i in range(args.num_flow_steps)
])
flow = flows.Flow(transform, distribution).to(device)
