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