def get_classifier(model, parameter_dim, observation_dim): if model == "linear": classifier = nn.Linear(parameter_dim + observation_dim, 1) elif model == "mlp": hidden_dim = 50 classifier = nn.Sequential( nn.Linear(parameter_dim + observation_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim), nn.BatchNorm1d(hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 1), ) elif model == "resnet": classifier = nets.ResidualNet( in_features=parameter_dim + observation_dim, out_features=1, hidden_features=50, context_features=None, num_blocks=2, activation=F.relu, dropout_probability=0.0, use_batch_norm=False, ) else: raise ValueError(f"'model' must be one of ['linear', 'mlp', 'resnet'].") return classifier
def make_scalar_classifier(dim, type="resnet", blocks=5, hidden_units=128, use_batch_norm=True): logger.info( f"Making scalar {type} classifier with {blocks} blocks and {hidden_units} hidden units each" ) if type == "resnet": classifier = nn.Sequential( nets.ResidualNet( in_features=dim, out_features=1, hidden_features=hidden_units, num_blocks=blocks, dropout_probability=0.0, use_batch_norm=use_batch_norm, ), nn.Sigmoid(), ) elif type == "mlp": assert not use_batch_norm classifier = nn.Sequential( nets.MLP( in_shape=(dim, ), out_shape=(1, ), hidden_sizes=[hidden_units for _ in range(blocks)], ), nn.Sigmoid(), ) else: raise ValueError(f"Unknown scalar classifier type {type}") return classifier
def transform_net_factory(in_features, out_features): return nets.ResidualNet( in_features=in_features, out_features=out_features, hidden_features=hidden_features, num_blocks=num_transform_blocks, dropout_probability=dropout_prob, use_batch_norm=use_batch_norm, )
def create_resnet(in_features, out_features): return nets.ResidualNet( in_features, out_features, hidden_features=hidden_features, num_blocks=num_blocks_per_layer, activation=activation, dropout_probability=dropout_probability, use_batch_norm=batch_norm_within_layers, )
def create_base_transform(i, param_dim, context_dim=None, hidden_dim=512, num_transform_blocks=2, activation='relu', dropout_probability=0.0, batch_norm=False, num_bins=8, tail_bound=1., apply_unconditional_transform=False, base_transform_type='rq-coupling' ): """Build a base NSF transform of x, conditioned on y. This uses the PiecewiseRationalQuadraticCoupling transform or the MaskedPiecewiseRationalQuadraticAutoregressiveTransform, as described in the Neural Spline Flow paper (https://arxiv.org/abs/1906.04032). Code is adapted from the uci.py example from https://github.com/bayesiains/nsf. A coupling flow fixes half the components of x, and applies a transform to the remaining components, conditioned on the fixed components. This is a restricted form of an autoregressive transform, with a single split into fixed/transformed components. The transform here is a neural spline flow, where the flow is parametrized by a residual neural network that depends on x_fixed and y. The residual network consists of a sequence of two-layer fully-connected blocks. Arguments: i {int} -- index of transform in sequence param_dim {int} -- dimensionality of x Keyword Arguments: context_dim {int} -- dimensionality of y (default: {None}) hidden_dim {int} -- number of hidden units per layer (default: {512}) num_transform_blocks {int} -- number of transform blocks comprising the transform (default: {2}) activation {str} -- activation function (default: {'relu'}) dropout_probability {float} -- probability of dropping out a unit (default: {0.0}) batch_norm {bool} -- whether to use batch normalization (default: {False}) num_bins {int} -- number of bins for the spline (default: {8}) tail_bound {[type]} -- [description] (default: {1.}) apply_unconditional_transform {bool} -- whether to apply an unconditional transform to fixed components (default: {False}) base_transform_type {str} -- type of base transform ([rq-coupling], rq-autoregressive) Returns: Transform -- the NSF transform """ if activation == 'elu': activation_fn = F.elu elif activation == 'relu': activation_fn = F.relu elif activation == 'leaky_relu': activation_fn = F.leaky_relu else: activation_fn = F.relu # Default print('Invalid activation function specified. Using ReLU.') if base_transform_type == 'rq-coupling': return transforms.PiecewiseRationalQuadraticCouplingTransform( mask=utils.create_alternating_binary_mask( param_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=hidden_dim, context_features=context_dim, num_blocks=num_transform_blocks, activation=activation_fn, dropout_probability=dropout_probability, use_batch_norm=batch_norm ) ), num_bins=num_bins, tails='linear', tail_bound=tail_bound, apply_unconditional_transform=apply_unconditional_transform ) elif base_transform_type == 'rq-autoregressive': return transforms.MaskedPiecewiseRationalQuadraticAutoregressiveTransform( features=param_dim, hidden_features=hidden_dim, context_features=context_dim, num_bins=num_bins, tails='linear', tail_bound=tail_bound, num_blocks=num_transform_blocks, use_residual_blocks=True, random_mask=False, activation=activation_fn, dropout_probability=dropout_probability, use_batch_norm=batch_norm ) else: raise ValueError
def create_net(in_features, out_features): return nets.ResidualNet(in_features, out_features, hidden_features=30, num_blocks=5)
def get_neural_posterior(model, parameter_dim, observation_dim, simulator): # Everything is a flow because we need to normalize parameters based on prior. mean, std = simulator.normalization_parameters normalizing_transform = AffineTransform(shift=-mean / std, scale=1 / std) if model == "mdn": hidden_features = 50 neural_posterior = MultivariateGaussianMDN( features=parameter_dim, context_features=observation_dim, hidden_features=hidden_features, hidden_net=nn.Sequential( nn.Linear(observation_dim, hidden_features), nn.ReLU(), nn.Dropout(p=0.0), nn.Linear(hidden_features, hidden_features), nn.ReLU(), nn.Linear(hidden_features, hidden_features), nn.ReLU(), ), num_components=20, custom_initialization=True, ) elif model == "made": num_mixture_components = 5 transform = normalizing_transform distribution = MADEMoG( features=parameter_dim, hidden_features=50, context_features=observation_dim, num_blocks=2, num_mixture_components=num_mixture_components, use_residual_blocks=True, random_mask=False, activation=F.relu, dropout_probability=0.0, use_batch_norm=False, custom_initialization=True, ) neural_posterior = Flow(transform, distribution) elif model == "maf": transform = CompositeTransform( [ CompositeTransform( [ transforms.MaskedAffineAutoregressiveTransform( features=parameter_dim, hidden_features=50, context_features=observation_dim, num_blocks=2, use_residual_blocks=False, random_mask=False, activation=F.tanh, dropout_probability=0.0, use_batch_norm=True, ), transforms.RandomPermutation(features=parameter_dim), ] ) for _ in range(5) ] ) transform = CompositeTransform([normalizing_transform, transform,]) distribution = StandardNormal((parameter_dim,)) neural_posterior = Flow(transform, distribution) elif model == "nsf": transform = CompositeTransform( [ CompositeTransform( [ transforms.PiecewiseRationalQuadraticCouplingTransform( mask=create_alternating_binary_mask( features=parameter_dim, even=(i % 2 == 0) ), transform_net_create_fn=lambda in_features, out_features: nets.ResidualNet( in_features=in_features, out_features=out_features, hidden_features=50, context_features=observation_dim, num_blocks=2, activation=F.relu, dropout_probability=0.0, use_batch_norm=False, ), num_bins=10, tails="linear", tail_bound=3.0, apply_unconditional_transform=False, ), transforms.LULinear(parameter_dim, identity_init=True), ] ) for i in range(5) ] ) distribution = StandardNormal((parameter_dim,)) neural_posterior = Flow(transform, distribution) else: raise ValueError return neural_posterior
def get_neural_likelihood(model, parameter_dim, observation_dim): if model == "mdn": hidden_features = 50 neural_likelihood = MultivariateGaussianMDN( features=observation_dim, context_features=parameter_dim, hidden_features=hidden_features, hidden_net=nn.Sequential( nn.Linear(parameter_dim, hidden_features), nn.BatchNorm1d(hidden_features), nn.ReLU(), nn.Dropout(p=0.0), nn.Linear(hidden_features, hidden_features), nn.BatchNorm1d(hidden_features), nn.ReLU(), nn.Linear(hidden_features, hidden_features), nn.BatchNorm1d(hidden_features), nn.ReLU(), ), num_components=20, custom_initialization=True, ) elif model == "made": neural_likelihood = MixtureOfGaussiansMADE( features=observation_dim, hidden_features=50, context_features=parameter_dim, num_blocks=4, num_mixture_components=10, use_residual_blocks=True, random_mask=False, activation=F.relu, use_batch_norm=True, dropout_probability=0.0, custom_initialization=True, ) elif model == "maf": transform = CompositeTransform( [ CompositeTransform( [ MaskedAffineAutoregressiveTransform( features=observation_dim, hidden_features=50, context_features=parameter_dim, num_blocks=2, use_residual_blocks=False, random_mask=False, activation=F.tanh, dropout_probability=0.0, use_batch_norm=True, ), RandomPermutation(features=observation_dim), ] ) for _ in range(5) ] ) distribution = StandardNormal((observation_dim,)) neural_likelihood = Flow(transform, distribution) elif model == "nsf": transform = CompositeTransform( [ CompositeTransform( [ PiecewiseRationalQuadraticCouplingTransform( mask=create_alternating_binary_mask( features=observation_dim, even=(i % 2 == 0) ), transform_net_create_fn=lambda in_features, out_features: nets.ResidualNet( in_features=in_features, out_features=out_features, hidden_features=50, context_features=parameter_dim, num_blocks=2, activation=F.relu, dropout_probability=0.0, use_batch_norm=False, ), num_bins=10, tails="linear", tail_bound=3.0, apply_unconditional_transform=False, ), LULinear(observation_dim, identity_init=True), ] ) for i in range(5) ] ) distribution = StandardNormal((observation_dim,)) neural_likelihood = Flow(transform, distribution) else: raise ValueError return neural_likelihood
def get_flow( model: str, dim_distribution: int, dim_context: Optional[int] = None, embedding: Optional[torch.nn.Module] = None, hidden_features: int = 50, made_num_mixture_components: int = 10, made_num_blocks: int = 4, flow_num_transforms: int = 5, mean=0.0, std=1.0, ) -> torch.nn.Module: """Density estimator Args: model: Model, one of maf / made / nsf dim_distribution: Dim of distribution dim_context: Dim of context embedding: Embedding network hidden_features: For all, number of hidden features made_num_mixture_components: For MADEs only, number of mixture components made_num_blocks: For MADEs only, number of blocks flow_num_transforms: For flows only, number of transforms mean: For normalization std: For normalization Returns: Neural network """ standardizing_transform = transforms.AffineTransform(shift=-mean / std, scale=1 / std) features = dim_distribution context_features = dim_context if model == "made": transform = standardizing_transform distribution = distributions_.MADEMoG( features=features, hidden_features=hidden_features, context_features=context_features, num_blocks=made_num_blocks, num_mixture_components=made_num_mixture_components, use_residual_blocks=True, random_mask=False, activation=torch.relu, dropout_probability=0.0, use_batch_norm=False, custom_initialization=True, ) neural_net = flows.Flow(transform, distribution, embedding) elif model == "maf": transform = transforms.CompositeTransform([ transforms.CompositeTransform([ transforms.MaskedAffineAutoregressiveTransform( features=features, hidden_features=hidden_features, context_features=context_features, num_blocks=2, use_residual_blocks=False, random_mask=False, activation=torch.tanh, dropout_probability=0.0, use_batch_norm=True, ), transforms.RandomPermutation(features=features), ]) for _ in range(flow_num_transforms) ]) transform = transforms.CompositeTransform( [standardizing_transform, transform]) distribution = distributions_.StandardNormal((features, )) neural_net = flows.Flow(transform, distribution, embedding) elif model == "nsf": transform = transforms.CompositeTransform([ transforms.CompositeTransform([ transforms.PiecewiseRationalQuadraticCouplingTransform( mask=create_alternating_binary_mask(features=features, even=(i % 2 == 0)), transform_net_create_fn=lambda in_features, out_features: nets.ResidualNet( in_features=in_features, out_features=out_features, hidden_features=hidden_features, context_features=context_features, num_blocks=2, activation=torch.relu, dropout_probability=0.0, use_batch_norm=False, ), num_bins=10, tails="linear", tail_bound=3.0, apply_unconditional_transform=False, ), transforms.LULinear(features, identity_init=True), ]) for i in range(flow_num_transforms) ]) transform = transforms.CompositeTransform( [standardizing_transform, transform]) distribution = distributions_.StandardNormal((features, )) neural_net = flows.Flow(transform, distribution, embedding) elif model == "nsf_bounded": transform = transforms.CompositeTransform([ transforms.CompositeTransform([ transforms.PiecewiseRationalQuadraticCouplingTransform( mask=create_alternating_binary_mask( features=dim_distribution, even=(i % 2 == 0)), transform_net_create_fn=lambda in_features, out_features: nets.ResidualNet( in_features=in_features, out_features=out_features, hidden_features=hidden_features, context_features=context_features, num_blocks=2, activation=F.relu, dropout_probability=0.0, use_batch_norm=False, ), num_bins=10, tails="linear", tail_bound=np.sqrt( 3), # uniform with sqrt(3) bounds has unit-variance apply_unconditional_transform=False, ), transforms.RandomPermutation(features=dim_distribution), ]) for i in range(flow_num_transforms) ]) transform = transforms.CompositeTransform( [standardizing_transform, transform]) distribution = StandardUniform(shape=(dim_distribution, )) neural_net = flows.Flow(transform, distribution, embedding) else: raise ValueError return neural_net