nflows依赖库架构及函数功能详解
- distributions
- flows
- nn
- transforms
- utils
- other new functions
distributions
distributions.base
class Distribution(nn.Module):Base class for all distribution objects.
可拓展函数:
def forward(self, *args):
def log_prob(self, inputs, context=None):
def _log_prob(self, inputs, context):
def sample(self, num_samples, context=None, batch_size=None):
def _sample(self, num_samples, context):
def sample_and_log_prob(self, num_samples, context=None):
def mean(self, context=None):
def _mean(self, context):
其中带下划线_的函数可重构。.
distributions.discrete
class ConditionalIndependentBernoulli(Distribution):An independent Bernoulli whose parameters are functions of a context.
def _compute_params(self, context):
def _log_prob(self, inputs, context):
def _sample(self, num_samples, context):
def _mean(self, context):
distributions.mixture
class MADEMoG(Distribution):MoG = MixtureOfGaussians 详情见from nflows.nn.nde import MixtureOfGaussiansMADE
函数声明
def __init__(
self,
features,
hidden_features,
context_features,
num_blocks=2,
num_mixture_components=1,
use_residual_blocks=True,
random_mask=False,
activation=F.relu,
dropout_probability=0.0,
use_batch_norm=False,
custom_initialization=False,
):
super().__init__()
self._made = MixtureOfGaussiansMADE(
features=features,
hidden_features=hidden_features,
context_features=context_features,
num_blocks=num_blocks,
num_mixture_components=num_mixture_components,
use_residual_blocks=use_residual_blocks,
random_mask=random_mask,
activation=activation,
dropout_probability=dropout_probability,
use_batch_norm=use_batch_norm,
custom_initialization=custom_initialization,
)
def _log_prob(self, inputs, context=None):
def _sample(self, num_samples, context):
distributions.normal
class StandardNormal(Distribution):A multivariate Normal with zero mean and unit covariance. N ( 0 , 1 ) \mathcal{N}(0,1) N(0,1)
def _log_prob(self, inputs, context=None):
def _sample(self, num_samples, context):
def _mean(self, context):
class ConditionalDiagonalNormal(Distribution):A diagonal multivariate Normal whose parameters are functions of a context.
def _compute_params(self, context):
def _log_prob(self, inputs, context=None):
def _sample(self, num_samples, context):
def _mean(self, context):
class DiagonalNormal(Distribution):A diagonal multivariate Normal with trainable parameters.
def __init__(self, shape):
"""Constructor.
Args:
shape: list, tuple or torch.Size, the shape of the input variables.
context_encoder: callable or None, encodes the context to the distribution parameters.
If None, defaults to the identity function.
"""
super().__init__()
self._shape = torch.Size(shape)
self.mean_ = nn.Parameter(torch.zeros(shape).reshape(1, -1))
self.log_std_ = nn.Parameter(torch.zeros(shape).reshape(1, -1))
self.register_buffer("_log_z",
torch.tensor(0.5 * np.prod(shape) * np.log(2 * np.pi),
dtype=torch.float64),
persistent=False)
def _log_prob(self, inputs, context=None):
def _sample(self, num_samples, context):
def _mean(self, context):
distributions.uniform
class BoxUniform(distributions.Independent):Multidimensional uniform distribution defined on a box. 主要用于更改 log_prob() 结果的形状。可以将创建与多元正态分布形状相同的对角正态分布，使得它们能够互换。
def __init__(
self,
low: Union[torch.Tensor, float],
high: Union[torch.Tensor, float],
reinterpreted_batch_ndims: int = 1,
):
"""Multidimensionqal uniform distribution defined on a box.
A `Uniform` distribution initialized with e.g. a parameter vector low or high of length 3 will
result in a /batch/ dimension of length 3.
A log_prob evaluation will then output three numbers, one for each of the independent Uniforms in the batch.
Instead, a `BoxUniform` initialized in the same way has three /event/ dimensions,
and returns a scalar log_prob corresponding to whether the evaluated point is in the box defined
by low and high or outside.
Refer to torch.distributions.Uniform and torch.distributions.Independent for further documentation.
Args:
low (Tensor or float): lower range (inclusive).
high (Tensor or float): upper range (exclusive).
reinterpreted_batch_ndims (int): the number of batch dims to reinterpret as event dims.
"""
super().__init__(
distributions.Uniform(low=low, high=high), reinterpreted_batch_ndims
)
class MG1Uniform(distributions.Uniform):no discription
class MG1Uniform(distributions.Uniform):
def log_prob(self, value):
return super().log_prob(self._to_noise(value))
def sample(self, sample_shape=torch.Size()):
return self._to_parameters(super().sample(sample_shape))
def _to_parameters(self, noise):
A_inv = torch.tensor([[1.0, 1, 0], [0, 1, 0], [0, 0, 1]])
return noise @ A_inv
def _to_noise(self, parameters):
A = torch.tensor([[1.0, -1, 0], [0, 1, 0], [0, 0, 1]])
return parameters @ A
class LotkaVolterraOscillating:1925年前后，美国数学家 Alfred J. Lotka (1880-1949) 和意大利数学家 Vito Volterra (1860-1940) 前后独立提出了一个类似的可以描述这种涨落规律的模型。该模型后来就因此以这两位研究者的名字命名，现被称为 Lotka-Volterra 模型，也常被称作猎食者-猎物（predator-prey）模型，后来逐渐发展成为了生态系统研究中的一个标志性的模型。【oscillating n. 振荡；摆动】
class LotkaVolterraOscillating:
def __init__(self):
mean = torch.log(torch.tensor([0.01, 0.5, 1, 0.01]))
sigma = 0.5
covariance = sigma ** 2 * torch.eye(4)
self._gaussian = distributions.MultivariateNormal(
loc=mean, covariance_matrix=covariance
)
self._uniform = BoxUniform(low=-5 * torch.ones(4), high=2 * torch.ones(4))
self._log_normalizer = -torch.log(
torch.erf((2 - mean) / sigma) - torch.erf((-5 - mean) / sigma)
).sum()
def log_prob(self, value):
unnormalized_log_prob = self._gaussian.log_prob(value) + self._uniform.log_prob(
value
)
return self._log_normalizer + unnormalized_log_prob
def sample(self, sample_shape=torch.Size()):
num_remaining_samples = sample_shape[0]
samples = []
while num_remaining_samples > 0:
candidate_samples = self._gaussian.sample((num_remaining_samples,))
uniform_log_prob = self._uniform.log_prob(candidate_samples)
accepted_samples = candidate_samples[~torch.isinf(uniform_log_prob)]
samples.append(accepted_samples.detach())
num_accepted = (~torch.isinf(uniform_log_prob)).sum().item()
num_remaining_samples -= num_accepted
# Aggregate collected samples.
samples = torch.cat(samples)
# Make sure we have the right amount.
samples = samples[: sample_shape[0], ...]
assert samples.shape[0] == sample_shape[0]
return samples
flows
flows.autoregressive
class MaskedAutoregressiveFlow(Flow):An autoregressive flow that uses affine transforms with masking. Masked Autoregressive Flow for Density Estimation, 2017NIPS
class MaskedAutoregressiveFlow(Flow):
def __init__(self):
mean = torch.log(torch.tensor([0.01, 0.5, 1, 0.01]))
sigma = 0.5
covariance = sigma ** 2 * torch.eye(4)
self._gaussian = distributions.MultivariateNormal(
loc=mean, covariance_matrix=covariance
)
self._uniform = BoxUniform(low=-5 * torch.ones(4), high=2 * torch.ones(4))
self._log_normalizer = -torch.log(
torch.erf((2 - mean) / sigma) - torch.erf((-5 - mean) / sigma)
).sum()
def log_prob(self, value):
unnormalized_log_prob = self._gaussian.log_prob(value) + self._uniform.log_prob(
value
)
return self._log_normalizer + unnormalized_log_prob
def sample(self, sample_shape=torch.Size()):
num_remaining_samples = sample_shape[0]
samples = []
while num_remaining_samples > 0:
candidate_samples = self._gaussian.sample((num_remaining_samples,))
uniform_log_prob = self._uniform.log_prob(candidate_samples)
accepted_samples = candidate_samples[~torch.isinf(uniform_log_prob)]
samples.append(accepted_samples.detach())
num_accepted = (~torch.isinf(uniform_log_prob)).sum().item()
num_remaining_samples -= num_accepted
# Aggregate collected samples.
samples = torch.cat(samples)
# Make sure we have the right amount.
samples = samples[: sample_shape[0], ...]
assert samples.shape[0] == sample_shape[0]
return samples
flows.base
class Flow(Distribution):Base class for all flow objects.
def __init__(self, transform, distribution, embedding_net=None):
Args:
transform: A `Transform` object, it transforms data into noise.
distribution: A `Distribution` object, the base distribution of the flow that
generates the noise.
embedding_net: A `nn.Module` which has trainable parameters to encode the
context (condition). It is trained jointly with the flow.
def _log_prob(self, inputs, context):
def _sample(self, num_samples, context):
def sample_and_log_prob(self, num_samples, context=None):
Generates samples from the flow, together with their log probabilities. For flows, this is more efficient that calling sample and log_prob separately.
def transform_to_noise(self, inputs, context=None):
Transforms given data into noise. Useful for goodness-of-fit checking.
Args:
inputs: A `Tensor` of shape [batch_size, ...], the data to be transformed.
context: A `Tensor` of shape [batch_size, ...] or None, optional context associated
with the data.
Returns:
A `Tensor` of shape [batch_size, ...], the noise.
flows.realnvp
class SimpleRealNVP(Flow):An simplified version of Real NVP for 1-dim inputs. This implementation uses 1-dim checkerboard masking but doesn’t use multi-scaling.Density estimation using Real NVP, 2017ICLR
class SimpleRealNVP(Flow):
def __init__(
self,
features,
hidden_features,
num_layers,
num_blocks_per_layer,
use_volume_preserving=False,
activation=F.relu,
dropout_probability=0.0,
batch_norm_within_layers=False,
batch_norm_between_layers=False,
):
if use_volume_preserving:
coupling_constructor = AdditiveCouplingTransform
else:
coupling_constructor = AffineCouplingTransform
mask = torch.ones(features)
mask[::2] = -1
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,
)
layers = []
for _ in range(num_layers):
transform = coupling_constructor(
mask=mask, transform_net_create_fn=create_resnet
)
layers.append(transform)
mask *= -1
if batch_norm_between_layers:
layers.append(BatchNorm(features=features))
super().__init__(
transform=CompositeTransform(layers),
distribution=StandardNormal([features]),
)
nn
nn.nde
nn.nde.made
class MaskedLinear(nn.Linear):A linear module with a masked weight matrix.
@classmethod
def _get_mask_and_degrees(cls, in_degrees, out_features, autoregressive_features, random_mask, is_output):
def forward(self, x):
return F.linear(x, self.weight * self.mask, self.bias)
class MaskedFeedforwardBlock(nn.Module)A feedforward block based on a masked linear module. In this implementation, the number of output features is taken to be equal to the number of input features.
self.linear = MaskedLinear(
in_degrees=in_degrees,
out_features=features,
autoregressive_features=autoregressive_features,
random_mask=random_mask,
is_output=False,
)
def forward(self, inputs, context=None):
if self.batch_norm:
outputs = self.batch_norm(inputs)
else:
outputs = inputs
outputs = self.linear(outputs)
outputs = self.activation(outputs)
outputs = self.dropout(outputs)
return outputs
class MaskedResidualBlock(nn.Module)A residual block containing masked linear modules.
# Masked linear.
linear_0 = MaskedLinear(
in_degrees=in_degrees,
out_features=features,
autoregressive_features=autoregressive_features,
random_mask=False,
is_output=False,
)
linear_1 = MaskedLinear(
in_degrees=linear_0.degrees,
out_features=features,
autoregressive_features=autoregressive_features,
random_mask=False,
is_output=False,
)
self.linear_layers = nn.ModuleList([linear_0, linear_1])
def forward(self, inputs, context=None):
class MADE(nn.Module)It can use either feedforward blocks or residual blocks (default is residual). Optionally, it can use batch norm or dropout within blocks (default is no).
self.initial_layer = MaskedLinear()
self.context_layer = nn.Linear(context_features, hidden_features)
for _ in range(num_blocks):
blocks.append(
block_constructor(
in_degrees=prev_out_degrees,
autoregressive_features=features,
context_features=context_features,
random_mask=random_mask,
activation=activation,
dropout_probability=dropout_probability,
use_batch_norm=use_batch_norm,
zero_initialization=True,
)
)
self.blocks = nn.ModuleList(blocks)
self.final_layer = MaskedLinear()
def forward(self, inputs, context=None):
temps = self.initial_layer(inputs)
if context is not None:
temps += self.context_layer(context)
for block in self.blocks:
temps = block(temps, context)
outputs = self.final_layer(temps)
return outputs
class MixtureOfGaussiansMADE(MADE):It can use either feedforward blocks or residual blocks (default is residual). Optionally, it can use batch norm or dropout within blocks (default is no).
def __init__(
self,
features,
hidden_features,
context_features=None,
num_blocks=2,
num_mixture_components=5,
use_residual_blocks=True,
random_mask=False,
activation=F.relu,
dropout_probability=0.0,
use_batch_norm=False,
epsilon=1e-2,
custom_initialization=True,
)
def forward(self, inputs, context=None):
return super().forward(inputs, context=context)
def log_prob(self, inputs, context=None):
def sample(self, num_samples, context=None):
def _initialize(self):
# Initialize mixture coefficient logits to near zero so that mixture coefficients
# are approximately uniform.
nn.nets
nn.nets.mlp
class MLP(nn.Module):A standard multi-layer perceptron.
def __init__(
self,
in_shape,
out_shape,
hidden_sizes,
activation=F.relu,
activate_output=False,
):
"""
Args:
in_shape: tuple, list or torch.Size, the shape of the input.
out_shape: tuple, list or torch.Size, the shape of the output.
hidden_sizes: iterable of ints, the hidden-layer sizes.
activation: callable, the activation function.
activate_output: bool, whether to apply the activation to the output.
"""
def forward(self, inputs):
nn.nets.resnet
class ResidualBlock(nn.Module):A general-purpose residual block. Works only with 1-dim inputs.
def __init__(
self,
in_shape,
out_shape,
hidden_sizes,
activation=F.relu,
activate_output=False,
):
"""
Args:
in_shape: tuple, list or torch.Size, the shape of the input.
out_shape: tuple, list or torch.Size, the shape of the output.
hidden_sizes: iterable of ints, the hidden-layer sizes.
activation: callable, the activation function.
activate_output: bool, whether to apply the activation to the output.
"""
def forward(self, inputs):
class ResidualNet(nn.Module):A general-purpose residual network. Works only with 1-dim inputs.
def __init__(
self,
in_features,
out_features,
hidden_features,
context_features=None,
num_blocks=2,
activation=F.relu,
dropout_probability=0.0,
use_batch_norm=False,
):
self.blocks = nn.ModuleList(
[
ResidualBlock(
features=hidden_features,
context_features=context_features,
activation=activation,
dropout_probability=dropout_probability,
use_batch_norm=use_batch_norm,
)
for _ in range(num_blocks)
]
def forward(self, inputs, context=None):
class ConvResidualBlock(nn.Module):
def __init__(
self,
in_features,
out_features,
hidden_features,
context_features=None,
num_blocks=2,
activation=F.relu,
dropout_probability=0.0,
use_batch_norm=False,
):
self.conv_layers = nn.ModuleList(
[nn.Conv2d(channels, channels, kernel_size=3, padding=1) for _ in range(2)]
)
def forward(self, inputs, context=None):
class ConvResidualNet(nn.Module):
def __init__(
self,
in_features,
out_features,
hidden_features,
context_features=None,
num_blocks=2,
activation=F.relu,
dropout_probability=0.0,
use_batch_norm=False,
):
self.blocks = nn.ModuleList(
[
ConvResidualBlock(
channels=hidden_channels,
context_channels=context_channels,
activation=activation,
dropout_probability=dropout_probability,
use_batch_norm=use_batch_norm,
)
for _ in range(num_blocks)
]
)
def forward(self, inputs, context=None):
transforms
transforms.splines
transforms.splines.cubic
unconstrained_cubic_splinecubic_spline
transforms.splines.linear
unconstrained_linear_splinelinear_spline
transforms.splines.quadratic
unconstrained_quadratic_splinequadratic_spline
transforms.splines.rational_quadratic
unconstrained_rational_quadratic_splinerational_quadratic_spline
transforms.autoregressive
class AutoregressiveTransform(Transform):Transforms each input variable with an invertible elementwise transformation. The parameters of each invertible elementwise transformation can be functions of previous input variables, but they must not depend on the current or any following input variables.
NOTE: Calculating the inverse transform is D times slower than calculating the forward transform, where D is the dimensionality of the input to the transform.
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
def _output_dim_multiplier(self):
def _elementwise_forward(self, inputs, autoregressive_params):
def _elementwise_inverse(self, inputs, autoregressive_params):
class MaskedAffineAutoregressiveTransform(AutoregressiveTransform):
made = made_module.MADE(
features=features,
hidden_features=hidden_features,
context_features=context_features,
num_blocks=num_blocks,
output_multiplier=self._output_dim_multiplier(),
use_residual_blocks=use_residual_blocks,
random_mask=random_mask,
activation=activation,
dropout_probability=dropout_probability,
use_batch_norm=use_batch_norm,
)
def _output_dim_multiplier(self):
def _elementwise_forward(self, inputs, autoregressive_params):
def _elementwise_inverse(self, inputs, autoregressive_params):
def _unconstrained_scale_and_shift(self, autoregressive_params):
class MaskedPiecewiseLinearAutoregressiveTransform(AutoregressiveTransform):
made = made_module.MADE(
features=features,
hidden_features=hidden_features,
context_features=context_features,
num_blocks=num_blocks,
output_multiplier=self._output_dim_multiplier(),
use_residual_blocks=use_residual_blocks,
random_mask=random_mask,
activation=activation,
dropout_probability=dropout_probability,
use_batch_norm=use_batch_norm,
)
def _output_dim_multiplier(self):
def _elementwise(self, inputs, autoregressive_params, inverse=False):
def _elementwise_forward(self, inputs, autoregressive_params):
def _elementwise_inverse(self, inputs, autoregressive_params):
class MaskedPiecewiseQuadraticAutoregressiveTransform(AutoregressiveTransform):
made = made_module.MADE(
features=features,
hidden_features=hidden_features,
context_features=context_features,
num_blocks=num_blocks,
output_multiplier=self._output_dim_multiplier(),
use_residual_blocks=use_residual_blocks,
random_mask=random_mask,
activation=activation,
dropout_probability=dropout_probability,
use_batch_norm=use_batch_norm,
)
def _output_dim_multiplier(self):
def _elementwise(self, inputs, autoregressive_params, inverse=False):
def _elementwise_forward(self, inputs, autoregressive_params):
def _elementwise_inverse(self, inputs, autoregressive_params):
class MaskedPiecewiseCubicAutoregressiveTransform(AutoregressiveTransform):
made = made_module.MADE(
features=features,
hidden_features=hidden_features,
context_features=context_features,
num_blocks=num_blocks,
output_multiplier=self._output_dim_multiplier(),
use_residual_blocks=use_residual_blocks,
random_mask=random_mask,
activation=activation,
dropout_probability=dropout_probability,
use_batch_norm=use_batch_norm,
)
def _output_dim_multiplier(self):
def _elementwise(self, inputs, autoregressive_params, inverse=False):
def _elementwise_forward(self, inputs, autoregressive_params):
def _elementwise_inverse(self, inputs, autoregressive_params):
class MaskedPiecewiseRationalQuadraticAutoregressiveTransform(AutoregressiveTransform):
made = made_module.MADE(
features=features,
hidden_features=hidden_features,
context_features=context_features,
num_blocks=num_blocks,
output_multiplier=self._output_dim_multiplier(),
use_residual_blocks=use_residual_blocks,
random_mask=random_mask,
activation=activation,
dropout_probability=dropout_probability,
use_batch_norm=use_batch_norm,
)
def _output_dim_multiplier(self):
def _elementwise(self, inputs, autoregressive_params, inverse=False):
def _elementwise_forward(self, inputs, autoregressive_params):
def _elementwise_inverse(self, inputs, autoregressive_params):
transforms.base
class Transform(nn.Module):Base class for all transform objects.
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class CompositeTransform(Transform):Composes several transforms into one, in the order they are given.
@staticmethod
def _cascade(inputs, funcs, context):
batch_size = inputs.shape[0]
outputs = inputs
total_logabsdet = inputs.new_zeros(batch_size)
for func in funcs:
outputs, logabsdet = func(outputs, context)
total_logabsdet += logabsdet
return outputs, total_logabsdet
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class MultiscaleCompositeTransform(Transform):A multiscale composite transform as described in the RealNVP paper. Splits the outputs along the given dimension after every transform, outputs one half, and passes the other half to further transforms. No splitting is done before the last transform. Note: Inputs could be of arbitrary shape, but outputs will always be flattened.
def add_transform(self, transform, transform_output_shape):
"""Add a transform. Must be called exactly `num_transforms` times.
Parameters:
transform: the `Transform` object to be added.
transform_output_shape: tuple, shape of transform's outputs, excl. the first batch
dimension.
Returns:
Input shape for the next transform, or None if adding the last transform.
"""
assert len(self._transforms) <= self._num_transforms
if len(self._transforms) == self._num_transforms:
raise RuntimeError(
"Adding more than {} transforms is not allowed.".format(
self._num_transforms
)
)
if (self._split_dim - 1) >= len(transform_output_shape):
raise ValueError("No split_dim in output shape")
if transform_output_shape[self._split_dim - 1] < 2:
raise ValueError(
"Size of dimension {} must be at least 2.".format(self._split_dim)
)
self._transforms.append(transform)
if len(self._transforms) != self._num_transforms: # Unless last transform.
output_shape = list(transform_output_shape)
output_shape[self._split_dim - 1] = (
output_shape[self._split_dim - 1] + 1
) // 2
output_shape = tuple(output_shape)
hidden_shape = list(transform_output_shape)
hidden_shape[self._split_dim - 1] = hidden_shape[self._split_dim - 1] // 2
hidden_shape = tuple(hidden_shape)
else:
# No splitting for last transform.
output_shape = transform_output_shape
hidden_shape = None
self._output_shapes.append(output_shape)
return hidden_shape
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class InverseTransform(Transform):Creates a transform that is the inverse of a given transform.
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
transforms.conv
class OneByOneConvolution(LULinear):An invertible 1x1 convolution with a fixed permutation, as introduced in the Glow paper.
def _lu_forward_inverse(self, inputs, inverse=False):
b, c, h, w = inputs.shape
inputs = inputs.permute(0, 2, 3, 1).reshape(b * h * w, c)
if inverse:
outputs, logabsdet = super().inverse(inputs)
else:
outputs, logabsdet = super().forward(inputs)
outputs = outputs.reshape(b, h, w, c).permute(0, 3, 1, 2)
logabsdet = logabsdet.reshape(b, h, w)
return outputs, torchutils.sum_except_batch(logabsdet)
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
transforms.coupling
class CouplingTransform(Transform):A base class for coupling layers. Supports 2D inputs (NxD), as well as 4D inputs for
images (NxCxHxW). For images the splitting is done on the channel dimension, using the provided 1D mask.
@property
def num_identity_features(self):
return len(self.identity_features)
@property
def num_transform_features(self):
return len(self.transform_features)
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
def _transform_dim_multiplier(self):
def _coupling_transform_forward(self, inputs, transform_params):
def _coupling_transform_inverse(self, inputs, transform_params):
class AffineCouplingTransform(CouplingTransform):An affine coupling layer that scales and shifts part of the variables. RealNVP
def _transform_dim_multiplier(self):
def _scale_and_shift(self, transform_params):
unconstrained_scale = transform_params[:, self.num_transform_features :, ...]
shift = transform_params[:, : self.num_transform_features, ...]
# scale = (F.softplus(unconstrained_scale) + 1e-3).clamp(0, 3)
scale = torch.sigmoid(unconstrained_scale + 2) + 1e-3
return scale, shift
def _coupling_transform_forward(self, inputs, transform_params):
def _coupling_transform_inverse(self, inputs, transform_params):
class AdditiveCouplingTransform(AffineCouplingTransform):An additive coupling layer, i.e. an affine coupling layer without scaling. NICE
def _transform_dim_multiplier(self):
def _scale_and_shift(self, transform_params):
unconstrained_scale = transform_params[:, self.num_transform_features :, ...]
shift = transform_params[:, : self.num_transform_features, ...]
# scale = (F.softplus(unconstrained_scale) + 1e-3).clamp(0, 3)
scale = torch.sigmoid(unconstrained_scale + 2) + 1e-3
return scale, shift
class AdditiveCouplingTransform(AffineCouplingTransform):An additive coupling layer, i.e. an affine coupling layer without scaling. NICE
def _transform_dim_multiplier(self):
def _scale_and_shift(self, transform_params):
unconstrained_scale = transform_params[:, self.num_transform_features :, ...]
shift = transform_params[:, : self.num_transform_features, ...]
# scale = (F.softplus(unconstrained_scale) + 1e-3).clamp(0, 3)
scale = torch.sigmoid(unconstrained_scale + 2) + 1e-3
return scale, shift
class PiecewiseCouplingTransform(CouplingTransform):(https://arxiv.org/abs/1410.8516)
def _coupling_transform_forward(self, inputs, transform_params):
return self._coupling_transform(inputs, transform_params, inverse=False)
def _coupling_transform_inverse(self, inputs, transform_params):
return self._coupling_transform(inputs, transform_params, inverse=True)
def _coupling_transform(self, inputs, transform_params, inverse=False):
if inputs.dim() == 4:
b, c, h, w = inputs.shape
# For images, reshape transform_params from Bx(C*?)xHxW to BxCxHxWx?
transform_params = transform_params.reshape(b, c, -1, h, w).permute(
0, 1, 3, 4, 2
)
elif inputs.dim() == 2:
b, d = inputs.shape
# For 2D data, reshape transform_params from Bx(D*?) to BxDx?
transform_params = transform_params.reshape(b, d, -1)
outputs, logabsdet = self._piecewise_cdf(inputs, transform_params, inverse)
return outputs, torchutils.sum_except_batch(logabsdet)
def _piecewise_cdf(self, inputs, transform_params, inverse=False):
class PiecewiseLinearCouplingTransform(PiecewiseCouplingTransform):Neural Importance Sampling
def _transform_dim_multiplier(self):
def _piecewise_cdf(self, inputs, transform_params, inverse=False):
class PiecewiseQuadraticCouplingTransform(PiecewiseCouplingTransform):Neural Importance Sampling
def _transform_dim_multiplier(self):
def _piecewise_cdf(self, inputs, transform_params, inverse=False):
class PiecewiseCubicCouplingTransform(PiecewiseCouplingTransform):
def _transform_dim_multiplier(self):
def _piecewise_cdf(self, inputs, transform_params, inverse=False):
class PiecewiseRationalQuadraticCouplingTransform(PiecewiseCouplingTransform):
def _transform_dim_multiplier(self):
def _piecewise_cdf(self, inputs, transform_params, inverse=False):
transforms.linear
class LinearCache:
"""Helper class to store the cache of a linear transform.
The cache consists of: the weight matrix, its inverse and its log absolute determinant.
"""
def __init__(self):
self.weight = None
self.inverse = None
self.logabsdet = None
def invalidate(self):
self.weight = None
self.inverse = None
self.logabsdet = None
class Linear(Transform):Abstract base class for linear transforms that parameterize a weight matrix.
def forward(self, inputs, context=None):
def _check_forward_cache(self):
def inverse(self, inputs, context=None):
def _check_inverse_cache(self):
def forward(self, inputs, context=None):
def train(self, mode=True):
def use_cache(self, mode=True):
def weight_and_logabsdet(self):
def weight_inverse_and_logabsdet(self):
def forward_no_cache(self, inputs):
def inverse_no_cache(self, inputs):
def weight(self):
"""Returns the weight matrix."""
def weight_inverse(self):
"""Returns the inverse weight matrix."""
def logabsdet(self):
"""Returns the log absolute determinant of the weight matrix."""
class NaiveLinear(Linear):A general linear transform that uses an unconstrained weight matrix.This transform explicitly computes the log absolute determinant in the forward direction and uses a linear solver in the inverse direction. Both forward and inverse directions have a cost of O(D^3), where D is the dimension of the input.
def forward_no_cache(self, inputs):
def inverse_no_cache(self, inputs):
def weight(self):
"""Returns the weight matrix."""
def weight_inverse(self):
"""Returns the inverse weight matrix."""
def weight_inverse_and_logabsdet(self):
def logabsdet(self):
"""Returns the log absolute determinant of the weight matrix."""
transforms.lu
class LULinear(Linear):A linear transform where we parameterize the LU decomposition of the weights.
def _initialize(self, identity_init):
def _create_lower_upper(self):
def forward_no_cache(self, inputs):
def inverse_no_cache(self, inputs):
def weight_inverse_and_logabsdet(self):
def weight(self):
def weight_inverse(self):
@property
def upper_diag(self):
return F.softplus(self.unconstrained_upper_diag) + self.eps
def logabsdet(self):
transforms.made
class MaskedLinear(nn.Linear)A linear module with a masked weight matrix.
@classmethod
def _get_mask_and_degrees(
cls, in_degrees, out_features, autoregressive_features, random_mask, is_output
):
def forward(self, x):
return F.linear(x, self.weight * self.mask, self.bias)
class MaskedFeedforwardBlock(nn.Module):A feedforward block based on a masked linear module. In this implementation, the number of output features is taken to be equal to the number of input features.
def __init__(
self,
in_degrees,
autoregressive_features,
context_features=None,
random_mask=False,
activation=F.relu,
dropout_probability=0.0,
use_batch_norm=False,
):
self.linear = MaskedLinear(
in_degrees=in_degrees,
out_features=features,
autoregressive_features=autoregressive_features,
random_mask=random_mask,
is_output=False,
)
def forward(self, inputs, context=None):
class MaskedResidualBlock(nn.Module):A residual block containing masked linear modules.
def __init__(
self,
in_degrees,
autoregressive_features,
context_features=None,
random_mask=False,
activation=F.relu,
dropout_probability=0.0,
use_batch_norm=False,
zero_initialization=True,
):
# Masked linear.
linear_0 = MaskedLinear(
in_degrees=in_degrees,
out_features=features,
autoregressive_features=autoregressive_features,
random_mask=False,
is_output=False,
)
linear_1 = MaskedLinear(
in_degrees=linear_0.degrees,
out_features=features,
autoregressive_features=autoregressive_features,
random_mask=False,
is_output=False,
)
self.linear_layers = nn.ModuleList([linear_0, linear_1])
def forward(self, inputs, context=None):
class MADE(nn.Module):It can use either feedforward blocks or residual blocks (default is residual). Optionally, it can use batch norm or dropout within blocks (default is no).
def __init__(
self,
features,
hidden_features,
context_features=None,
num_blocks=2,
output_multiplier=1,
use_residual_blocks=True,
random_mask=False,
activation=F.relu,
dropout_probability=0.0,
use_batch_norm=False,
):
def forward(self, inputs, context=None):
transforms.nonlinearities
class Tanh(Transform):
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class LogTanh(Transform):
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class LeakyReLU(Transform):
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class Sigmoid(Transform):
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
-
class Logit(InverseTransform): -
class GatedLinearUnit(Transform):
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class CauchyCDF(Transform):
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
-
class CauchyCDFInverse(InverseTransform): -
class CompositeCDFTransform(CompositeTransform): -
class PiecewiseLinearCDF(Transform):
def _spline(self, inputs, inverse=False):
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class PiecewiseCubicCDF(Transform):
def _spline(self, inputs, inverse=False):
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class PiecewiseRationalQuadraticCDF(Transform):
def _spline(self, inputs, inverse=False):
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
transforms.normalization
class BatchNorm(Transform):Transform that performs batch normalization.
Limitations:
* It works only for 1-dim inputs.
* Inverse is not available in training mode, only in eval mode.
def _spline(self, inputs, inverse=False):
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class ActNorm(Transform):Transform that performs activation normalization. Works for 2D and 4D inputs. For 4D inputs (images) normalization is performed per-channel, assuming BxCxHxW input shape.
@property
def scale(self):
return torch.exp(self.log_scale)
def _broadcastable_scale_shift(self, inputs):
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
def _initialize(self, inputs):
transforms.orthogonal
class HouseholderSequence(Transform):A sequence of Householder transforms. This class can be used as a way of parameterizing an
def __init__(self, features, num_transforms):
"""Constructor.
Args:
features: int, dimensionality of the input.
num_transforms: int, number of Householder transforms to use.
Raises:
TypeError: if arguments are not the right type.
"""
@staticmethod
def _apply_transforms(inputs, q_vectors):
"""Apply the sequence of transforms parameterized by given q_vectors to inputs.
Costs O(KDN), where:
- K is number of transforms
- D is dimensionality of inputs
- N is number of inputs
Args:
inputs: Tensor of shape [N, D]
q_vectors: Tensor of shape [K, D]
Returns:
A tuple of:
- A Tensor of shape [N, D], the outputs.
- A Tensor of shape [N], the log absolute determinants of the total transform.
"""
squared_norms = torch.sum(q_vectors ** 2, dim=-1)
outputs = inputs
for q_vector, squared_norm in zip(q_vectors, squared_norms):
temp = outputs @ q_vector # Inner product.
temp = torch.ger(temp, (2.0 / squared_norm) * q_vector) # Outer product.
outputs = outputs - temp
batch_size = inputs.shape[0]
logabsdet = inputs.new_zeros(batch_size)
return outputs, logabsdet
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
def matrix(self):
transforms.permutations
class Permutation(Transform):Permutes inputs on a given dimension using a given permutation.
@property
def _inverse_permutation(self):
return torch.argsort(self._permutation)
@staticmethod
def _permute(inputs, permutation, dim):
if dim >= inputs.ndimension():
raise ValueError("No dimension {} in inputs.".format(dim))
if inputs.shape[dim] != len(permutation):
raise ValueError(
"Dimension {} in inputs must be of size {}.".format(
dim, len(permutation)
)
)
batch_size = inputs.shape[0]
outputs = torch.index_select(inputs, dim, permutation)
logabsdet = inputs.new_zeros(batch_size)
return outputs, logabsdet
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class RandomPermutation(Permutation):Permutes using a random, but fixed, permutation. Only works with 1D inputs.class ReversePermutation(Permutation):Reverses the elements of the input. Only works with 1D inputs.
transforms.qr
class QRLinear(Linear):A linear module using the QR decomposition for the weight matrix.
def _initialize(self, identity_init):
def _create_upper(self):
def forward_no_cache(self, inputs):
def inverse_no_cache(self, inputs):
def weight(self):
def weight_inverse(self):
def logabsdet(self):
transforms.reshape
class SqueezeTransform(Transform):A transformation defined for image data that trades spatial dimensions for channel dimensions, i.e. “squeezes” the inputs along the channel dimensions.
def get_output_shape(self, c, h, w):
return (c * self.factor * self.factor, h // self.factor, w // self.factor)
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
transforms.standard
class IdentityTransform(Transform):Transform that leaves input unchanged.
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
class PointwiseAffineTransform(Transform):Forward transform X = X * scale + shift.
def _log_scale(self) -> Tensor:
return torch.log(self._scale)
def _batch_logabsdet(self, batch_shape: Iterable[int]) -> Tensor:
"""Return log abs det with input batch shape."""
if self._log_scale.numel() > 1:
return self._log_scale.expand(batch_shape).sum()
else:
# When log_scale is a scalar, we use n*log_scale, which is more
# numerically accurate than \sum_1^n log_scale.
return self._log_scale * torch.Size(batch_shape).numel()
def forward(self, inputs, context=None):
def inverse(self, inputs, context=None):
AffineTransform(PointwiseAffineTransform):
transforms.svd
class SVDLinear(Linear):A linear module using the SVD decomposition for the weight matrix.
@property
def diagonal(self):
return self.eps + F.softplus(self.unconstrained_diagonal)
@property
def log_diagonal(self):
return torch.log(self.diagonal)
def _initialize(self):
def forward_no_cache(self, inputs):
def inverse_no_cache(self, inputs):
def weight(self):
def weight_inverse(self):
def logabsdet(self):
utils
utils.torchutils
def tile(x, n):
def sum_except_batch(x, num_batch_dims=1):
def split_leading_dim(x, shape):
def merge_leading_dims(x, num_dims):
def repeat_rows(x, num_reps):
def tensor2numpy(x):
def logabsdet(x):
def random_orthogonal(size):
def get_num_parameters(model):
def create_alternating_binary_mask(features, even=True):
def create_mid_split_binary_mask(features):
def create_random_binary_mask(features):
def searchsorted(bin_locations, inputs, eps=1e-6):
def cbrt(x):
def get_temperature(max_value, bound=1 - 1e-3):
def gaussian_kde_log_eval(samples, query):
utils.typechecks
"""Functions that check types."""
def is_bool(x):
return isinstance(x, bool)
def is_int(x):
return isinstance(x, int)
def is_positive_int(x):
return is_int(x) and x > 0
def is_nonnegative_int(x):
return is_int(x) and x >= 0
def is_power_of_two(n):
if is_positive_int(n):
return not n & (n - 1)
else:
return False
other new functions
register_buffer()
661
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
