def conditional_normalizing_flow_factory(flow_depth, problem_dim, c_net_depth, c_net_h_dim, context_dim,
context_n_h_dim, context_n_depth, rich_context_dim, cuda):
# We define the base distribution
if cuda:
base_dist = dist.Normal(torch.zeros(problem_dim).cuda(), torch.ones(problem_dim).cuda())
else:
base_dist = dist.Normal(torch.zeros(problem_dim), torch.ones(problem_dim))
# We define the transformations
transforms = [conditional_affine_coupling(problem_dim,
context_dim,
hidden_dims=[c_net_h_dim for i in range(c_net_depth)], # Note array here to create multiple layers in DenseNN
rich_context_dim=rich_context_dim,
context_hidden_dims=[context_n_h_dim for i in range(context_n_depth)]) for i in range(flow_depth)]
# need a fix for this
perms = [permute(2, torch.tensor([1, 0])) for i in range(flow_depth)]
# We sandwich the AffineCouplings and permutes together. Unelegant hotfix to remove last permute but it works
flows = list(itertools.chain(*zip(transforms, perms)))[:-1]
# We define the normalizing flow wrapper
normalizing_flow = ConditionalNormalizingFlowWrapper(transforms, flows, base_dist)
if cuda:
normalizing_flow.cuda()
return normalizing_flow
def conditional_normalizing_flow_factory3(flow_depth, problem_dim, c_net_depth, c_net_h_dim, context_dim,
context_n_h_dim, context_n_depth, rich_context_dim, batchnorm_momentum, cuda,
coupling_dropout=None, context_dropout=None):
if cuda:
base_dist = dist.Normal(torch.zeros(problem_dim).cuda(), torch.ones(problem_dim).cuda())
else:
base_dist = dist.Normal(torch.zeros(problem_dim), torch.ones(problem_dim))
# We define the transformations
transforms = [conditional_affine_coupling2(input_dim=problem_dim,
context_dim=context_dim,
hidden_dims=[c_net_h_dim for i in range(c_net_depth)], # Note array here to create multiple layers in DenseNN
rich_context_dim=rich_context_dim,
dropout=coupling_dropout)
for i in range(flow_depth)]
# Permutes are needed to be able to transform all dimensions.
# Note that the transform is fixed here since we only have 2 dimensions. For more dimensions don't fix it and let it be random.
perms = [permute(input_dim=problem_dim, permutation=torch.tensor([1, 0])) for i in range(flow_depth)]
# If we want batchnorm add those in. Then sandwich the steps together to a flow
if batchnorm_momentum is None:
batchnorms = None
flows = list(itertools.chain(*zip(transforms, perms)))[:-1]
else:
batchnorms = [batchnorm(input_dim=problem_dim, momentum=batchnorm_momentum) for i in range(flow_depth)]
for bn in batchnorms:
bn.gamma.data += torch.ones(problem_dim)
flows = list(itertools.chain(*zip(batchnorms, transforms, perms)))[1:-1]
# We define the conditioning network
context_hidden_dims = [context_n_h_dim for i in range(context_n_depth)]
if context_dropout is None:
condinet = DenseNN(input_dim=context_dim, hidden_dims=context_hidden_dims, param_dims=[rich_context_dim])
else:
condinet = DropoutDenseNN(input_dim=context_dim, hidden_dims=context_hidden_dims, param_dims=[rich_context_dim],
dropout=context_dropout)
# We define the normalizing flow wrapper
normalizing_flow = ConditionalNormalizingFlowWrapper3(transforms, flows, base_dist, condinet, batchnorms)
if cuda:
normalizing_flow.cuda()
return normalizing_flow
def planar_normalizing_flow_factory(flow_depth, problem_dim, c_net_depth,
c_net_h_dim, batchnorm_momentum, cuda):
# We define the base distribution
if cuda:
base_dist = dist.Normal(
torch.zeros(problem_dim).cuda(),
torch.ones(problem_dim).cuda())
else:
base_dist = dist.Normal(torch.zeros(problem_dim),
torch.ones(problem_dim))
# We define the transformations
transforms = [
affine_coupling(input_dim=problem_dim,
hidden_dims=[c_net_h_dim for i in range(c_net_depth)])
for i in range(flow_depth)
]
# We need to permute dimensions to affect them both THIS NEEDS A FIX
perms = [permute(2, torch.tensor([1, 0])) for i in range(flow_depth)]
# If we want batchnorm add those in. Then sandwich the steps together to a flow
if batchnorm_momentum is None:
batchnorms = None
flows = list(itertools.chain(*zip(transforms, perms)))[:-1]
else:
batchnorms = [
batchnorm(input_dim=problem_dim, momentum=batchnorm_momentum)
for i in range(flow_depth)
]
for bn in batchnorms:
bn.gamma.data += torch.ones(problem_dim)
flows = list(
itertools.chain(*zip(batchnorms, transforms, perms)))[1:-1]
# We define the normalizing flow wrapper
normalizing_flow = NormalizingFlowWrapper(transforms, flows, base_dist,
batchnorms)
if cuda:
normalizing_flow.cuda()
return normalizing_flow
def __init__(self,
input_dim=2,
split_dim=1,
context_dim=1,
hidden_dim=128,
num_layers=1,
flow_length=10,
use_cuda=False):
super(ConditionalNormalizingFlow, self).__init__()
self.base_dist = dist.Normal(
torch.zeros(input_dim),
torch.ones(input_dim)) # base distribution is Isotropic Gaussian
self.param_dims = [input_dim - split_dim, input_dim - split_dim]
# Define series of bijective transformations
self.transforms = [
ConditionalAffineCoupling(
split_dim,
ConditionalDenseNN(split_dim, context_dim,
[hidden_dim] * num_layers, self.param_dims))
for _ in range(flow_length)
]
self.perms = [
permute(2, torch.tensor([1, 0])) for _ in range(flow_length)
]
self.bns = [BatchNorm(input_dim=1) for _ in range(flow_length)]
# Concatenate AffineCoupling layers with Permute and BatchNorm Layers
self.generative_flows = list(
itertools.chain(*zip(self.transforms, self.bns, self.perms))
)[:-2] # generative direction (z-->x)
self.normalizing_flows = self.generative_flows[::
-1] # normalizing direction (x-->z)
self.use_cuda = use_cuda
if self.use_cuda:
self.cuda()
nn.ModuleList(self.transforms).cuda()
self.base_dist = dist.Normal(
torch.zeros(input_dim).cuda(),
torch.ones(input_dim).cuda())
def test_permute_shapes(self):
for shape in [(3, ), (3, 4), (3, 4, 2)]:
input_dim = shape[-1]
self._test_shape(shape, T.permute(input_dim))
def test_permute_inverses(self):
for input_dim in [2, 5, 10]:
self._test_inverse(input_dim, T.permute(input_dim))
def combi_conditional_normalizing_flow_factory(flow_depth,
problem_dim,
c_net_depth,
c_net_h_dim,
context_dim,
context_n_h_dim,
context_n_depth,
rich_context_dim,
batchnorm_momentum,
cuda,
coupling_dropout=None,
context_dropout=None,
planar_first=True):
assert (
flow_depth % 2 == 0
), "The flow depth must be divisible by 2 to allow both Planar and AC transforms"
if cuda:
base_dist = dist.Normal(
torch.zeros(problem_dim).cuda(),
torch.ones(problem_dim).cuda())
else:
base_dist = dist.Normal(torch.zeros(problem_dim),
torch.ones(problem_dim))
# We define the transformations
affine_couplings = [
conditional_affine_coupling2(
input_dim=problem_dim,
context_dim=context_dim,
hidden_dims=[c_net_h_dim for i in range(c_net_depth)],
rich_context_dim=rich_context_dim,
dropout=coupling_dropout) for i in range(flow_depth // 2)
]
planars = [
inverted_conditional_planar(
input_dim=problem_dim,
context_dim=rich_context_dim,
hidden_dims=[c_net_h_dim for i in range(c_net_depth)],
) for i in range(flow_depth // 2)
]
transforms = affine_couplings + planars
# Permutes are needed to be able to transform all dimensions.
# Note that the transform is fixed here since we only have 2 dimensions. For more dimensions let it be random.
perms = [
permute(input_dim=problem_dim, permutation=torch.tensor([1, 0]))
for i in range(flow_depth // 2)
]
# Assemble the flow
if planar_first:
flows = list(
itertools.chain(*zip(planars, affine_couplings, perms)))[:-1]
else:
flows = list(
itertools.chain(*zip(affine_couplings, planars, perms)))[:-1]
# If we want batchnorm add those in. Then sandwich the steps together to a flow
if batchnorm_momentum is None:
batchnorms = None
else:
bn_flow = flows[:1]
batchnorms = []
for trans in flows[1:]:
if isinstance(trans, ConditionalAffineCoupling2) or isinstance(
trans, InvertedConditionalPlanar):
batchnorms.append(
batchnorm(input_dim=problem_dim,
momentum=batchnorm_momentum))
bn_flow.append(batchnorms[-1])
bn_flow.append(trans)
flows = bn_flow
for bn in batchnorms:
bn.gamma.data += torch.ones(problem_dim)
# We define the conditioning network
context_hidden_dims = [context_n_h_dim for i in range(context_n_depth)]
if context_dropout is None:
condinet = DenseNN(input_dim=context_dim,
hidden_dims=context_hidden_dims,
param_dims=[rich_context_dim])
else:
condinet = DropoutDenseNN(input_dim=context_dim,
hidden_dims=context_hidden_dims,
param_dims=[rich_context_dim],
dropout=context_dropout)
# We define the normalizing flow wrapper
normalizing_flow = ConditionalNormalizingFlowWrapper3(
transforms, flows, base_dist, condinet, batchnorms)
if cuda:
normalizing_flow.cuda()
return normalizing_flow