def build_resnet_classifier(
batch_x: Tensor,
batch_y: Tensor,
z_score_x: Optional[str] = "independent",
z_score_y: Optional[str] = "independent",
hidden_features: int = 50,
embedding_net_x: nn.Module = nn.Identity(),
embedding_net_y: nn.Module = nn.Identity(),
num_blocks: int = 2,
dropout_probability: float = 0.0,
use_batch_norm: bool = False,
) -> nn.Module:
"""Builds ResNet classifier.
In SNRE, the classifier will receive batches of thetas and xs.
Args:
batch_x: Batch of xs, used to infer dimensionality and (optional) z-scoring.
batch_y: Batch of ys, used to infer dimensionality and (optional) z-scoring.
z_score_x: Whether to z-score xs passing into the network, can be one of:
- `none`, or None: do not z-score.
- `independent`: z-score each dimension independently.
- `structured`: treat dimensions as related, therefore compute mean and std
over the entire batch, instead of per-dimension. Should be used when each
sample is, for example, a time series or an image.
z_score_y: Whether to z-score ys passing into the network, same options as
z_score_x.
hidden_features: Number of hidden features.
embedding_net_x: Optional embedding network for x.
embedding_net_y: Optional embedding network for y.
Returns:
Neural network.
"""
check_data_device(batch_x, batch_y)
check_embedding_net_device(embedding_net=embedding_net_x, datum=batch_y)
check_embedding_net_device(embedding_net=embedding_net_y, datum=batch_y)
# Infer the output dimensionalities of the embedding_net by making a forward pass.
x_numel = embedding_net_x(batch_x[:1]).numel()
y_numel = embedding_net_y(batch_y[:1]).numel()
neural_net = nets.ResidualNet(
in_features=x_numel + y_numel,
out_features=1,
hidden_features=hidden_features,
context_features=None,
num_blocks=num_blocks,
activation=relu,
dropout_probability=dropout_probability,
use_batch_norm=use_batch_norm,
)
input_layer = build_input_layer(
batch_x, batch_y, z_score_x, z_score_y, embedding_net_x, embedding_net_y
)
neural_net = nn.Sequential(input_layer, neural_net)
return neural_net
def classifier_nn(
model,
prior: torch.distributions.Distribution,
context: torch.Tensor,
hidden_features: int = 50,
) -> torch.nn.Module:
"""Neural classifier
Args:
model: Model, one of linear / mlp / resnet
prior: Prior distribution
context: Observation
hidden_features: For all, number of hidden features
Returns:
Neural network
"""
parameter_dim = prior.sample([1]).shape[1]
context = utils.torchutils.atleast_2d(context)
observation_dim = torch.tensor([context.shape[1:]])
if model == "linear":
neural_net = nn.Linear(parameter_dim + observation_dim, 1)
elif model == "mlp":
neural_net = nn.Sequential(
nn.Linear(parameter_dim + observation_dim, hidden_features),
nn.BatchNorm1d(hidden_features),
nn.ReLU(),
nn.Linear(hidden_features, hidden_features),
nn.BatchNorm1d(hidden_features),
nn.ReLU(),
nn.Linear(hidden_features, 1),
)
elif model == "resnet":
neural_net = nets.ResidualNet(
in_features=parameter_dim + observation_dim,
out_features=1,
hidden_features=hidden_features,
context_features=None,
num_blocks=2,
activation=torch.relu,
dropout_probability=0.0,
use_batch_norm=False,
)
else:
raise ValueError(f"'model' must be one of ['linear', 'mlp', 'resnet'].")
return neural_net
def build_resnet_classifier(
batch_x: Tensor = None,
batch_y: Tensor = None,
z_score_x: bool = True,
z_score_y: bool = True,
hidden_features: int = 50,
embedding_net_x: nn.Module = nn.Identity(),
embedding_net_y: nn.Module = nn.Identity(),
device: str = 'cpu'
) -> nn.Module:
"""Builds ResNet classifier.
In SNRE, the classifier will receive batches of thetas and xs.
Args:
batch_x: Batch of xs, used to infer dimensionality and (optional) z-scoring.
batch_y: Batch of ys, used to infer dimensionality and (optional) z-scoring.
z_score_x: Whether to z-score xs passing into the network.
z_score_y: Whether to z-score ys passing into the network.
hidden_features: Number of hidden features.
embedding_net_x: Optional embedding network for x.
embedding_net_y: Optional embedding network for y.
Returns:
Neural network.
"""
# Infer the output dimensionalities of the embedding_net by making a forward pass.
x_numel = embedding_net_x(batch_x[:1]).numel()
y_numel = embedding_net_y(batch_y[:1]).numel()
neural_net = nets.ResidualNet(
in_features=x_numel + y_numel,
out_features=1,
hidden_features=hidden_features,
context_features=None,
num_blocks=2,
activation=relu,
dropout_probability=0.0,
use_batch_norm=False,
device=device
)
input_layer = build_input_layer(
batch_x, batch_y, z_score_x, z_score_y, embedding_net_x, embedding_net_y
)
neural_net = nn.Sequential(input_layer, neural_net).to(device)
return neural_net
def build_nn(theta):
classifier = nets.ResidualNet(
in_features=theta.shape[1],
out_features=2,
hidden_features=hidden_features,
context_features=None,
num_blocks=num_blocks,
activation=relu,
dropout_probability=dropout_probability,
use_batch_norm=True,
)
input_layer = build_input_layer(theta, z_score_theta,
embedding_net_theta)
classifier = nn.Sequential(input_layer, classifier)
return classifier
def posterior_nn(
model: str,
prior: torch.distributions.Distribution,
context: torch.Tensor,
embedding: Optional[torch.nn.Module] = None,
hidden_features: int = 50,
mdn_num_components: int = 20,
made_num_mixture_components: int = 10,
made_num_blocks: int = 4,
flow_num_transforms: int = 5,
) -> torch.nn.Module:
"""Neural posterior density estimator
Args:
model: Model, one of maf / mdn / made / nsf
prior: Prior distribution
context: Observation
embedding: Embedding network
hidden_features: For all, number of hidden features
mdn_num_components: For MDNs only, number of components
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
Returns:
Neural network
"""
mean, std = (prior.mean, prior.stddev)
standardizing_transform = transforms.AffineTransform(
shift=-mean / std, scale=1 / std
)
parameter_dim = prior.sample([1]).shape[1]
context = utils.torchutils.atleast_2d(context)
observation_dim = torch.tensor([context.shape[1:]])
if model == "mdn":
neural_net = 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=mdn_num_components,
custom_initialization=True,
)
elif model == "made":
transform = standardizing_transform
distribution = distributions_.MADEMoG(
features=parameter_dim,
hidden_features=hidden_features,
context_features=observation_dim,
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=parameter_dim,
hidden_features=hidden_features,
context_features=observation_dim,
num_blocks=2,
use_residual_blocks=False,
random_mask=False,
activation=torch.tanh,
dropout_probability=0.0,
use_batch_norm=True,
),
transforms.RandomPermutation(features=parameter_dim),
]
)
for _ in range(flow_num_transforms)
]
)
transform = transforms.CompositeTransform([standardizing_transform, transform,])
distribution = distributions_.StandardNormal((parameter_dim,))
neural_net = flows.Flow(transform, distribution, embedding)
elif model == "nsf":
transform = transforms.CompositeTransform(
[
transforms.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=hidden_features,
context_features=observation_dim,
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(parameter_dim, identity_init=True),
]
)
for i in range(flow_num_transforms)
]
)
transform = transforms.CompositeTransform([standardizing_transform, transform,])
distribution = distributions_.StandardNormal((parameter_dim,))
neural_net = flows.Flow(transform, distribution, embedding)
else:
raise ValueError
return neural_net
def build_nsf(
batch_x: Tensor = None,
batch_y: Tensor = None,
z_score_x: bool = True,
z_score_y: bool = True,
hidden_features: int = 50,
num_transforms: int = 5,
num_bins: int = 10,
embedding_net: nn.Module = nn.Identity(),
**kwargs,
) -> nn.Module:
"""Builds NSF p(x|y).
Args:
batch_x: Batch of xs, used to infer dimensionality and (optional) z-scoring.
batch_y: Batch of ys, used to infer dimensionality and (optional) z-scoring.
z_score_x: Whether to z-score xs passing into the network.
z_score_y: Whether to z-score ys passing into the network.
hidden_features: Number of hidden features.
num_transforms: Number of transforms.
num_bins: Number of bins used for the splines.
embedding_net: Optional embedding network for y.
kwargs: Additional arguments that are passed by the build function but are not
relevant for maf and are therefore ignored.
Returns:
Neural network.
"""
x_numel = batch_x[0].numel()
# Infer the output dimensionality of the embedding_net by making a forward pass.
y_numel = embedding_net(batch_y[:1]).numel()
if x_numel == 1:
class ContextSplineMap(nn.Module):
"""
Neural network from `context` to the spline parameters.
We cannot use the resnet as conditioner to learn each dimension conditioned
on the other dimensions (because there is only one). Instead, we learn the
spline parameters directly. In the case of conditinal density estimation,
we make the spline parameters conditional on the context. This is
implemented in this class.
"""
def __init__(
self,
in_features: int,
out_features: int,
hidden_features: int,
context_features: int,
):
"""
Initialize neural network that learns to predict spline parameters.
Args:
in_features: Unused since there is no `conditioner` in 1D.
out_features: Number of spline parameters.
hidden_features: Number of hidden units.
context_features: Number of context features.
"""
super().__init__()
# `self.hidden_features` is only defined such that nflows can infer
# a scaling factor for initializations.
self.hidden_features = hidden_features
# Use a non-linearity because otherwise, there will be a linear
# mapping from context features onto distribution parameters.
self.spline_predictor = nn.Sequential(
nn.Linear(context_features, self.hidden_features),
nn.ReLU(),
nn.Linear(self.hidden_features, self.hidden_features),
nn.ReLU(),
nn.Linear(self.hidden_features, out_features),
)
def __call__(self, inputs: Tensor, context: Tensor, *args,
**kwargs) -> Tensor:
"""
Return parameters of the spline given the context.
Args:
inputs: Unused. It would usually be the other dimensions, but in
1D, there are no other dimensions.
context: Context features.
Returns:
Spline parameters.
"""
return self.spline_predictor(context)
mask_in_layer = lambda i: tensor([1], dtype=uint8)
conditioner = lambda in_features, out_features: ContextSplineMap(
in_features,
out_features,
hidden_features,
context_features=y_numel)
if num_transforms > 1:
warn(
f"You are using `num_transforms={num_transforms}`. When estimating a "
f"1D density, you will not get any performance increase by using "
f"multiple transforms with NSF. We recommend setting "
f"`num_transforms=1` for faster training (see also 'Change "
f"hyperparameters of density esitmators' here: "
f"https://www.mackelab.org/sbi/tutorial/04_density_estimators/)."
)
else:
mask_in_layer = lambda i: create_alternating_binary_mask(
features=x_numel, even=(i % 2 == 0))
conditioner = lambda in_features, out_features: nets.ResidualNet(
in_features=in_features,
out_features=out_features,
hidden_features=hidden_features,
context_features=y_numel,
num_blocks=2,
activation=relu,
dropout_probability=0.0,
use_batch_norm=False,
)
transform = transforms.CompositeTransform([
transforms.CompositeTransform([
transforms.PiecewiseRationalQuadraticCouplingTransform(
mask=mask_in_layer(i),
transform_net_create_fn=conditioner,
num_bins=num_bins,
tails="linear",
tail_bound=3.0,
apply_unconditional_transform=False,
),
transforms.LULinear(x_numel, identity_init=True),
]) for i in range(num_transforms)
])
if z_score_x:
transform_zx = standardizing_transform(batch_x)
transform = transforms.CompositeTransform([transform_zx, transform])
if z_score_y:
embedding_net = nn.Sequential(standardizing_net(batch_y),
embedding_net)
distribution = distributions_.StandardNormal((x_numel, ))
neural_net = flows.Flow(transform, distribution, embedding_net)
return neural_net
def build_nsf(
batch_x: Tensor = None,
batch_y: Tensor = None,
z_score_x: bool = True,
z_score_y: bool = True,
hidden_features: int = 50,
num_transforms: int = 5,
embedding_net: nn.Module = nn.Identity(),
**kwargs,
) -> nn.Module:
"""Builds NSF p(x|y).
Args:
batch_x: Batch of xs, used to infer dimensionality and (optional) z-scoring.
batch_y: Batch of ys, used to infer dimensionality and (optional) z-scoring.
z_score_x: Whether to z-score xs passing into the network.
z_score_y: Whether to z-score ys passing into the network.
hidden_features: Number of hidden features.
num_transforms: Number of transforms.
embedding_net: Optional embedding network for y.
kwargs: Additional arguments that are passed by the build function but are not
relevant for maf and are therefore ignored.
Returns:
Neural network.
"""
x_numel = batch_x[0].numel()
# Infer the output dimensionality of the embedding_net by making a forward pass.
y_numel = embedding_net(batch_y[:1]).numel()
if x_numel == 1:
raise NotImplementedError
transform = transforms.CompositeTransform(
[
transforms.CompositeTransform(
[
transforms.PiecewiseRationalQuadraticCouplingTransform(
mask=create_alternating_binary_mask(
features=x_numel, 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=y_numel,
num_blocks=num_blocks,
activation=relu,
dropout_probability=0.0,
use_batch_norm=False,
),
num_bins=10,
tails="linear",
tail_bound=3.0,
apply_unconditional_transform=False,
),
transforms.LULinear(x_numel, identity_init=True),
]
)
for i in range(num_transforms)
]
)
if z_score_x:
transform_zx = standardizing_transform(batch_x)
transform = transforms.CompositeTransform([transform_zx, transform])
if z_score_y:
embedding_net = nn.Sequential(standardizing_net(batch_y), embedding_net)
distribution = distributions_.StandardNormal((x_numel,))
neural_net = flows.Flow(transform, distribution, embedding_net)
return neural_net