def transform_factory(input_dim): transform = T.batchnorm(input_dim) transform._inverse( torch.normal(torch.arange(0., input_dim), torch.arange(1., 1. + input_dim) / input_dim)) transform.eval() return transform
def test_batchnorm_shapes(self): for shape in [(3,), (3, 4), (3, 4, 2)]: input_dim = shape[-1] transform = T.batchnorm(input_dim) transform._inverse(torch.normal(torch.arange(0., input_dim), torch.arange(1., 1. + input_dim) / input_dim)) transform.eval() self._test_shape(shape, transform)
def inverted_conditional_planar_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, 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 = [ 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) ] # If we want batchnorm add those in. Then sandwich the steps together to a flow if batchnorm_momentum is None: batchnorms = None flows = transforms 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)))[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 test_batchnorm_jacobians(self): for input_dim in [2, 5, 10]: transform = T.batchnorm(input_dim) transform._inverse( torch.normal(torch.arange(0., input_dim), torch.arange(1., 1. + input_dim) / input_dim)) transform.eval() self._test_jacobian(input_dim, transform)
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 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