def sample(self, num_samples, context=None, batch_size=None):
"""Generates samples from the distribution. Samples can be generated in batches.
Args:
num_samples: int, number of samples to generate.
context: Tensor or None, conditioning variables. If None, the context is ignored.
batch_size: int or None, number of samples per batch. If None, all samples are generated
in one batch.
Returns:
A Tensor containing the samples, with shape [num_samples, ...] if context is None, or
[context_size, num_samples, ...] if context is given.
"""
if not check.is_positive_int(num_samples):
raise TypeError("Number of samples must be a positive integer.")
if context is not None:
context = torch.as_tensor(context)
if batch_size is None:
return self._sample(num_samples, context)
else:
if not check.is_positive_int(batch_size):
raise TypeError("Batch size must be a positive integer.")
num_batches = num_samples // batch_size
num_leftover = num_samples % batch_size
samples = [
self._sample(batch_size, context) for _ in range(num_batches)
]
if num_leftover > 0:
samples.append(self._sample(num_leftover, context))
return torch.cat(samples, dim=0)
def repeat_rows(x, num_reps):
"""Each row of tensor `x` is repeated `num_reps` times along leading dimension."""
if not check.is_positive_int(num_reps):
raise TypeError("Number of repetitions must be a positive integer.")
shape = x.shape
x = x.unsqueeze(1)
x = x.expand(shape[0], num_reps, *shape[1:])
return merge_leading_dims(x, num_dims=2)
def tile(x, n):
if not check.is_positive_int(n):
raise TypeError("Argument 'n' must be a positive integer.")
x_ = x.reshape(-1)
x_ = x_.repeat(n)
x_ = x_.reshape(n, -1)
x_ = x_.transpose(1, 0)
x_ = x_.reshape(-1)
return x_
def __init__(self, permutation, dim=1):
if permutation.ndimension() != 1:
raise ValueError("Permutation must be a 1D tensor.")
if not check.is_positive_int(dim):
raise ValueError("dim must be a positive integer.")
super().__init__()
self._dim = dim
self.register_buffer("_permutation", permutation)
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.
"""
if not check.is_positive_int(features):
raise TypeError("Number of features must be a positive integer.")
if not check.is_positive_int(num_transforms):
raise TypeError("Number of transforms must be a positive integer.")
super().__init__()
self.features = features
self.num_transforms = num_transforms
# TODO: are randn good initial values?
# these vectors are orthogonal to the hyperplanes through which we reflect
# self.q_vectors = nets.Parameter(torch.randn(num_transforms, features))
# self.q_vectors = nets.Parameter(torch.eye(num_transforms // 2, features))
import numpy as np
def tile(a, dim, n_tile):
if a.nelement() == 0:
return a
init_dim = a.size(dim)
repeat_idx = [1] * a.dim()
repeat_idx[dim] = n_tile
a = a.repeat(*(repeat_idx))
order_index = torch.Tensor(
np.concatenate([
init_dim * np.arange(n_tile) + i for i in range(init_dim)
])).long()
return torch.index_select(a, dim, order_index)
qv = tile(torch.eye(num_transforms // 2, features), 0, 2)
if np.mod(num_transforms,
2) != 0: # odd number of transforms, including 1
qv = torch.cat((qv, torch.zeros(1, features)))
qv[-1, num_transforms // 2] = 1
self.q_vectors = nn.Parameter(qv)
def merge_leading_dims(x, num_dims):
"""Reshapes the tensor `x` such that the first `num_dims` dimensions are merged to one."""
if not check.is_positive_int(num_dims):
raise TypeError("Number of leading dims must be a positive integer.")
if num_dims > x.dim():
raise ValueError(
"Number of leading dims can't be greater than total number of dims."
)
new_shape = torch.Size([-1]) + x.shape[num_dims:]
return torch.reshape(x, new_shape)
def __init__(self, features, using_cache=False):
if not check.is_positive_int(features):
raise TypeError("Number of features must be a positive integer.")
super().__init__()
self.features = features
self.bias = nn.Parameter(torch.zeros(features))
# Caching flag and values.
self.using_cache = using_cache
self.cache = LinearCache()
def __init__(self, features, eps=1e-5, momentum=0.1, affine=True):
if not check.is_positive_int(features):
raise TypeError("Number of features must be a positive integer.")
super().__init__()
self.momentum = momentum
self.eps = eps
constant = np.log(np.exp(1 - eps) - 1)
self.unconstrained_weight = nn.Parameter(constant * torch.ones(features))
self.bias = nn.Parameter(torch.zeros(features))
self.register_buffer("running_mean", torch.zeros(features))
self.register_buffer("running_var", torch.zeros(features))
def __init__(self, permutation, dim=1):
if permutation.ndimension() != 1:
raise ValueError("Permutation must be a 1D tensor.")
if not check.is_positive_int(dim):
raise ValueError("dim must be a positive integer.")
super().__init__()
self._dim = dim
self.register_buffer("_permutation", permutation)
if self._permutation.get_device() == -1:
self.device = 'cpu'
else:
self.device = 'cuda:' + str(int(self._permutation.get_device()))
def __init__(self, features):
"""
Transform that performs activation normalization. Works for 2D and 4D inputs. For 4D
inputs (images) normalization is performed per-channel, assuming BxCxHxW input shape.
Reference:
> D. Kingma et. al., Glow: Generative flow with invertible 1x1 convolutions, NeurIPS 2018.
"""
if not check.is_positive_int(features):
raise TypeError("Number of features must be a positive integer.")
super().__init__()
self.initialized = False
self.log_scale = nn.Parameter(torch.zeros(features))
self.shift = nn.Parameter(torch.zeros(features))
def __init__(self, num_transforms, split_dim=1):
"""Constructor.
Args:
num_transforms: int, total number of transforms to be added.
split_dim: dimension along which to split.
"""
if not check.is_positive_int(split_dim):
raise TypeError("Split dimension must be a positive integer.")
super().__init__()
self._transforms = nn.ModuleList()
self._output_shapes = []
self._num_transforms = num_transforms
self._split_dim = split_dim
def __init__(self, features, dim=1):
if not check.is_positive_int(features):
raise ValueError("Number of features must be a positive integer.")
super().__init__(torch.arange(features - 1, -1, -1), dim)
def __init__(self, features, device, dim=1):
if not check.is_positive_int(features):
raise ValueError("Number of features must be a positive integer.")
super().__init__(torch.randperm(features).to(device), dim)
