def _check_repetition_indices(self, context: SamplingContext):
assert context.repetition_indices.shape[
0] == context.parent_indices.shape[0]
if self.num_repetitions > 1 and context.repetition_indices is None:
raise Exception(
f"Sum layer has multiple repetitions (num_repetitions=={self.num_repetitions}) but repetition_indices argument was None, expected a Long tensor size #samples."
)
if self.num_repetitions == 1 and context.repetition_indices is None:
context.repetition_indices = torch.zeros(context.n,
dtype=int,
device=self.__device)
def sample(self,
n: int = None,
context: SamplingContext = None) -> SamplingContext:
"""Method to sample from this layer, based on the parents output.
Args:
n (int): Number of instances to sample.
indices (torch.Tensor): Parent sampling output.
Returns:
torch.Tensor: Index into tensor which paths should be followed.
Output should be of size: in_features, out_channels.
"""
# If this is a root node
if context.is_root:
if self.num_repetitions == 1:
# If there is only a single repetition, create new sampling context
context.parent_indices = torch.zeros(context.n,
1,
dtype=int,
device=self.__device)
context.repetition_indices = torch.zeros(context.n,
dtype=int,
device=self.__device)
return context
else:
raise Exception(
"Cannot start sampling from Product layer with num_repetitions > 1 and no context given."
)
else:
# Repeat the parent indices, e.g. [0, 2, 3] -> [0, 0, 2, 2, 3, 3] depending on the cardinality
indices = torch.repeat_interleave(context.parent_indices,
repeats=self.cardinality,
dim=1)
# Remove padding
if self._pad:
indices = indices[:, :-self._pad]
context.parent_indices = indices
return context
def sample(self,
n: int = None,
context: SamplingContext = None) -> SamplingContext:
"""Method to sample from this layer, based on the parents output.
Args:
n: Number of samples.
indices (torch.Tensor): Parent sampling output
Returns:
torch.Tensor: Index into tensor which paths should be followed.
Output should be of size: in_features, out_channels.
"""
# If this is a root node
if context.is_root:
if self.num_repetitions == 1:
# If there is only a single repetition, create new sampling context
context.parent_indices = torch.zeros(context.n,
1,
dtype=int,
device=self.__device)
context.repetition_indices = torch.zeros(context.n,
dtype=int,
device=self.__device)
return context
else:
raise Exception(
"Cannot start sampling from CrossProduct layer with num_repetitions > 1 and no context given."
)
# Map flattened indexes back to coordinates to obtain the chosen input_channel for each feature
indices = self.unraveled_channel_indices[context.parent_indices]
indices = indices.view(indices.shape[0], -1)
# Remove padding
if self._pad:
indices = indices[:, :-self._pad]
context.parent_indices = indices
return context
def sample(self,
n: int = None,
class_index=None,
evidence: torch.Tensor = None):
"""
Sample from the distribution represented by this SPN.
Possible valid inputs:
- `n`: Generates `n` samples.
- `n` and `class_index (int)`: Generates `n` samples from P(X | C = class_index).
- `class_index (List[int])`: Generates `len(class_index)` samples. Each index `c_i` in `class_index` is mapped
to a sample from P(X | C = c_i)
- `evidence`: If evidence is given, samples conditionally and fill NaN values.
Args:
n: Number of samples to generate.
class_index: Class index. Can be either an int in combination with a value for `n` which will result in `n`
samples from P(X | C = class_index). Or can be a list of ints which will map each index `c_i` in the
list to a sample from P(X | C = c_i).
evidence: Evidence that can be provided to condition the samples. If evidence is given, `n` and
`class_index` must be `None`. Evidence must contain NaN values which will be imputed according to the
distribution represented by the SPN. The result will contain the evidence and replace all NaNs with the
sampled values.
Returns:
torch.Tensor: Samples generated according to the distribution specified by the SPN.
"""
assert class_index is None or evidence is None, "Cannot provide both, evidence and class indices."
assert n is None or evidence is None, "Cannot provide both, number of samples to generate (n) and evidence."
# Check if evidence contains nans
if evidence is not None:
assert (evidence != evidence).any(), "Evidence has no NaN values."
# Set n to the number of samples in the evidence
n = evidence.shape[0]
with provide_evidence(self, evidence): # May be None but that's ok
# If class is given, use it as base index
if class_index is not None:
if isinstance(class_index, list):
indices = torch.tensor(class_index,
device=self.__device).view(-1, 1)
n = indices.shape[0]
else:
indices = torch.empty(size=(n, 1), device=self.__device)
indices.fill_(class_index)
# Create new sampling context
ctx = SamplingContext(n=n,
parent_indices=indices,
repetition_indices=None)
else:
# Start sampling one of the C root nodes TODO: check what happens if C=1
ctx = self._sampling_root.sample(n=n)
# Sample from RatSpn root layer: Results are indices into the stacked output channels of all repetitions
ctx.repetition_indices = torch.zeros(n,
dtype=int,
device=self.__device)
ctx = self.root.sample(context=ctx)
# Indexes will now point to the stacked channels of all repetitions (R * S^2 (if D > 1)
# or R * I^2 (else)).
root_in_channels = self.root.in_channels // self.config.R
# Obtain repetition indices
ctx.repetition_indices = (ctx.parent_indices //
root_in_channels).squeeze(1)
# Shift indices
ctx.parent_indices = ctx.parent_indices % root_in_channels
# Now each sample in `indices` belongs to one repetition, index in `repetition_indices`
# Continue at layers
# Sample inner modules
for layer in reversed(self._inner_layers):
if isinstance(layer, Sum):
ctx = layer.sample(context=ctx)
elif isinstance(layer, CrossProduct):
ctx = layer.sample(context=ctx)
else:
raise Exception(
"Only Sum or CrossProduct is allowed as intermediate layer."
)
# Sample leaf
samples = self._leaf.sample(context=ctx)
# Invert permutation
for i in range(n):
rep_index = ctx.repetition_indices[i]
inv_rand_indices = invert_permutation(
self.rand_indices[:, rep_index])
samples[i, :] = samples[i, inv_rand_indices]
if evidence is not None:
# Update NaN entries in evidence with the sampled values
nan_indices = torch.isnan(evidence)
# First make a copy such that the original object is not changed
evidence = evidence.clone()
evidence[nan_indices] = samples[nan_indices]
return evidence
else:
return samples
def sample(self,
n: int = None,
context: SamplingContext = None) -> SamplingContext:
"""Method to sample from this layer, based on the parents output.
Output is always a vector of indices into the channels.
Args:
repetition_indices (List[int]): An index into the repetition axis for each sample.
Can be None if `num_repetitions==1`.
indices (torch.Tensor): Parent sampling output.
n (int): Number of samples.
Returns:
torch.Tensor: Index into tensor which paths should be followed.
"""
# Sum weights are of shape: [D, IC, OC, R]
# We now want to use `indices` to access one in_channel for each in_feature x out_channels block
# index is of size in_feature
weights = self.weights.data
d, ic, oc, r = weights.shape
n = context.n
# Create sampling context if this is a root layer
if context.is_root:
assert oc == 1 and r == 1, "Cannot start sampling from non-root layer."
# Initialize rep indices
context.repetition_indices = torch.zeros(n,
dtype=int,
device=self.__device)
# Select weights, repeat n times along the last dimension
weights = weights[:, :, [0] * n, 0] # Shape: [D, IC, N]
# Move sample dimension to the first axis: [feat, channels, batch] -> [batch, feat, channels]
weights = weights.permute(2, 0, 1) # Shape: [N, D, IC]
else:
# If this is not the root node, use the paths (out channels), specified by the parent layer
self._check_repetition_indices(context)
tmp = torch.zeros(n, d, ic, device=self.__device)
for i in range(n):
tmp[i, :, :] = weights[range(self.in_features), :,
context.parent_indices[i],
context.repetition_indices[i]]
weights = tmp
# Check dimensions
assert weights.shape == (n, d, ic)
# Apply softmax to ensure they are proper probabilities
log_weights = F.log_softmax(weights, dim=2)
# If evidence is given, adjust the weights with the likelihoods of the observed paths
if self._is_input_cache_enabled and self._input_cache is not None:
for i in range(n):
# Reweight the i-th samples weights by its likelihood values at the correct repetition
log_weights[i, :, :] += self._input_cache[
i, :, :, context.repetition_indices[i]]
# If sampling context is MPE, set max weight to 1 and rest to zero, such that the maximum index will be sampled
if context.is_mpe:
# Get index of largest weight along in-channel dimension
indices = log_weights.argmax(dim=2)
else:
# Create categorical distribution and use weights as logits.
#
# Use the Gumble-Softmax trick to obtain one-hot indices of the categorical distribution
# represented by the given logits. (Use Gumble-Softmax instead of Categorical
# to allow for gradients).
#
# The code below is an approximation of:
#
# >> dist = torch.distributions.Categorical(logits=log_weights)
# >> indices = dist.sample()
cats = torch.arange(ic, device=log_weights.device)
one_hot = F.gumbel_softmax(logits=log_weights, hard=True, dim=-1)
indices = (one_hot * cats).sum(-1).long()
# Update parent indices
context.parent_indices = indices
return context