def i_sample(self, shape=None, as_dist=False):
shape = size_getter(shape)
dist = TransformedDistribution(
self.noise0.expand(shape),
AffineTransform(self.i_mean(),
self.i_scale(),
event_dim=self._event_dim))
if as_dist:
return dist
return dist.sample()
def SampleAction(self, mean, std):
# mean and ln_var are predicted by the neural network, this function
mu = mean
sig = std * 0.3 # constraining the standard deviation to at maximum 0.3mu
u_range = self.args['U_UB'] - self.args['U_LB']
GPol = norm(
mu, sig
) # defining gaussian distribution with mean and std as parameterised
scale = AffineTransform(self.args['U_LB'], u_range)
GPol = TransformedDistribution(GPol, scale)
action = GPol.sample() # drawing randomly from normal distribution
assert len(action) == 1
logGP = GPol.log_prob(
action) # calculating log probability of action taken
return action.cpu(), logGP
def i_sample(self, shape=None, as_dist=False):
"""
Samples from the initial distribution.
:param shape: The number of samples
:type shape: int|tuple[int]
:param as_dist: Whether to return the new value as a distribution
:type as_dist: bool
:return: Samples from the initial distribution
:rtype: torch.Tensor|float
"""
shape = ((shape,) if isinstance(shape, int) else shape) or torch.Size([])
loc = concater(self.f0(*parameter_caster(len(shape), *self.theta)))
scale = concater(self.g0(*parameter_caster(len(shape), *self.theta)))
dist = TransformedDistribution(self.noise0.expand(shape), self._transform(loc, scale))
if as_dist:
return dist
return dist.sample()
def test_transformed_distribution(base_batch_dim, base_event_dim,
transform_dim, num_transforms, sample_shape):
shape = torch.Size([2, 3, 4, 5])
base_dist = Normal(0, 1)
base_dist = base_dist.expand(shape[4 - base_batch_dim - base_event_dim:])
if base_event_dim:
base_dist = Independent(base_dist, base_event_dim)
transforms = [
AffineTransform(torch.zeros(shape[4 - transform_dim:]), 1),
ReshapeTransform((4, 5), (20, )),
ReshapeTransform((3, 20), (6, 10))
]
transforms = transforms[:num_transforms]
transform = ComposeTransform(transforms)
# Check validation in .__init__().
if base_batch_dim + base_event_dim < transform.domain.event_dim:
with pytest.raises(ValueError):
TransformedDistribution(base_dist, transforms)
return
d = TransformedDistribution(base_dist, transforms)
# Check sampling is sufficiently expanded.
x = d.sample(sample_shape)
assert x.shape == sample_shape + d.batch_shape + d.event_shape
num_unique = len(set(x.reshape(-1).tolist()))
assert num_unique >= 0.9 * x.numel()
# Check log_prob shape on full samples.
log_prob = d.log_prob(x)
assert log_prob.shape == sample_shape + d.batch_shape
# Check log_prob shape on partial samples.
y = x
while y.dim() > len(d.event_shape):
y = y[0]
log_prob = d.log_prob(y)
assert log_prob.shape == d.batch_shape
def propagate(self, x, as_dist=False):
"""
Propagates the model forward conditional on the previous state and current parameters.
:param x: The previous state
:type x: torch.Tensor|float
:param as_dist: Whether to return the new value as a distribution
:type as_dist: bool
:return: Samples from the model
:rtype: torch.Tensor|float
"""
loc, scale = self.mean(x), self.scale(x)
if isinstance(self, Observable):
shape = _get_shape(loc if loc.dim() > scale.dim() else scale, self.ndim)
else:
shape = _get_shape(x, self.ndim)
dist = TransformedDistribution(self.noise.expand(shape), self._transform(loc, scale))
if as_dist:
return dist
return dist.sample()
class NICE(nn.Module):
def __init__(self, prior, coupling, in_out_dim, mid_dim, hidden,
bottleneck, compress, device, n_layers):
"""Initialize a NICE.
Args:
coupling: number of coupling layers.
in_out_dim: input/output dimensions.
mid_dim: number of units in a hidden layer.
hidden: number of hidden layers.
device: run on cpu or gpu
"""
super(NICE, self).__init__()
self.device = device
if prior == 'gaussian':
self.prior = torch.distributions.Normal(
torch.tensor(0.).to(device),
torch.tensor(1.).to(device))
elif prior == 'logistic':
self.prior = TransformedDistribution(
Uniform(
torch.tensor(0.).to(device),
torch.tensor(1.).to(device)),
[SigmoidTransform().inv,
AffineTransform(loc=0., scale=1.)])
else:
raise ValueError('Prior not implemented.')
self.in_out_dim = in_out_dim
self.coupling = coupling
self.n_layers = n_layers
layer = AdditiveCoupling if coupling == 'additive' else AffineCoupling
self.coupling_layers = nn.ModuleList([
layer(in_out_dim, mid_dim, hidden, i % 2)
for i in range(self.n_layers)
]).to(device)
self.scale = Scaling(in_out_dim).to(device)
self.bottleneck_factor = compress
self.bottleneck_loss = nn.MSELoss()
self.bottleneck = bottleneck
def f_inverse(self, z):
"""Transformation g: Z -> X (inverse of f).
Args:
z: tensor in latent space Z.
Returns:
transformed tensor in data space X.
"""
x, det = self.scale(z, reverse=True)
for layer in reversed(self.coupling_layers):
x, _ = layer(x, 0, reverse=True)
return x
def f(self, x):
"""Transformation f: X -> Z (inverse of g).
Args:
x: tensor in data space X.
Returns:
transformed tensor in latent space Z and log determinant Jacobian
"""
log_det_J = 0
for layer in self.coupling_layers:
x, log_det_J = layer(x, log_det_J)
z, det = self.scale(x)
return z, log_det_J + det
def loss(self, x):
"""Computes data log-likelihood.
(See Section 3.3 in the NICE paper.)
Args:
x: input minibatch.
Returns:
log-likelihood of input.
"""
z, log_det_J = self.f(x)
slices = [
z[:, i::self.bottleneck_factor]
for i in range(self.bottleneck_factor)
]
s = torch.stack(slices).permute(1, 0, 2)
if self.bottleneck == 'redundancy':
bottleneck_loss = torch.sum(torch.var(s, dim=1), dim=1)
log_ll = 0.0
for slice in slices:
log_ll += torch.sum(self.prior.log_prob(slice), dim=1)
if self.bottleneck == 'null':
winner = slices[-1]
loser = torch.ones_like(winner)
bottleneck_loss = 0.0
for slice in slices[:-1]:
bottleneck_loss += self.bottleneck_loss(slice, loser)
log_ll = torch.sum(self.prior.log_prob(winner), dim=1)
log_det_J -= np.log(
256
) * self.in_out_dim #/ self.bottleneck_factor #log det for rescaling from [0.256] (after dequantization) to [0,1]
#log_ll = torch.sum(self.prior.log_prob(z), dim=1)
return log_ll + log_det_J, bottleneck_loss
def sample(self, size):
"""Generates samples.
Args:
size: number of samples to generate.
Returns:
samples from the data space X.
"""
z = self.prior.sample(
(size, self.in_out_dim // self.bottleneck_factor)).to(self.device)
z_tag = torch.zeros((size, self.in_out_dim)).to(self.device)
if self.bottleneck == 'redundancy':
for i in range(self.bottleneck_factor):
z_tag[:, i::self.bottleneck_factor] = z
if self.bottleneck == 'null':
for i in range(self.bottleneck_factor - 1):
z_tag[:, i::self.bottleneck_factor] = torch.ones_like(z)
z_tag[:, self.bottleneck_factor - 1::self.bottleneck_factor] = z
return self.f_inverse(z_tag)
def forward(self, x):
"""Forward pass.
Args:
x: input minibatch.
Returns:
log-likelihood of input.
"""
x = x.to(self.device)
return self.loss(x)
