def test_sum_as_intermediate_node(self):
"""Check that sum node returns the correct sample indices when used as indermediate node."""
# Some values for the sum layer
in_features = 10
in_channels = 3
out_channels = 5
num_repetitions = 7
n = 2
parent_indices = torch.randint(out_channels, size=(n, in_features))
# Create sum layer
sum_layer = layers.Sum(
in_features=in_features, in_channels=in_channels, out_channels=out_channels, num_repetitions=num_repetitions
)
# Choose `in_features` number of random indexes from 0 to in_channels-1 which will have probability of 1.0 in
# the sum layer weight tensor
rand_indxs = torch.randint(in_channels, size=(in_features, num_repetitions))
rep_idxs = torch.randint(num_repetitions, size=(n,))
# Artificially set sum weights (probabilities) to 1.0
weights = torch.zeros(in_features, in_channels, out_channels, num_repetitions)
for r in range(num_repetitions):
weights[range(in_features), rand_indxs[:, r], :, r] = 1.0
sum_layer.weights = nn.Parameter(torch.log(weights))
# Perform sampling
ctx = SamplingContext(n=n, parent_indices=parent_indices, repetition_indices=rep_idxs)
sum_layer.sample(context=ctx)
# Assert that the sample indexes are those where the weights were set to 1.0
for i in range(n):
self.assertTrue((rand_indxs[:, rep_idxs[i]] == ctx.parent_indices[i, :]).all())
def test_normal_leaf(self):
# Setup leaf layer
out_channels = 10
in_features = 10
num_repetitions = 5
leaf = distributions.Normal(out_channels=out_channels, in_features=in_features, num_repetitions=num_repetitions)
# Set leaf layer mean to some random int
leaf.means.data = torch.randint(
low=-100, high=100, size=(1, in_features, out_channels, num_repetitions)
).float()
# Set leaf layer std to 0 such that the samples will all be the mean (so we can actually make assertions in the end)
leaf.stds.data = torch.zeros(size=(1, in_features, out_channels, num_repetitions)).float()
# Create some random indices into the out_channels axis
parent_indices = torch.randint(high=out_channels, size=(1, in_features,))
repetition_indices = torch.randint(high=num_repetitions, size=(1,))
# Perform sampling
ctx = SamplingContext(n=1, parent_indices=parent_indices, repetition_indices=repetition_indices)
result = leaf.sample(context=ctx)
# Expected sampling
expected_result = leaf.means.data[:, range(in_features), parent_indices, repetition_indices[0]]
# Run assertions
self.assertTrue(((result - expected_result).abs() < 1e-6).all())
def test_product_shape_as_root_node(self):
"""Check that the product node has the correct sampling shape when used as root."""
prod_layer = layers.Product(in_features=10,
cardinality=2,
num_repetitions=1)
ctx = SamplingContext(n=5)
ctx = prod_layer.sample(context=ctx)
self.assertTrue(ctx.parent_indices.shape[0] == 5)
self.assertTrue(ctx.parent_indices.shape[1] == 1)
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 None:
if self.num_repetitions == 1:
# If there is only a single repetition, create new sampling context
return SamplingContext(
n=n,
parent_indices=torch.zeros(n,
1,
dtype=int,
device=self.__device),
repetition_indices=torch.zeros(n,
dtype=int,
device=self.__device),
)
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 None:
if self.num_repetitions == 1:
# If there is only a single repetition, create new sampling context
return SamplingContext(
n=n,
parent_indices=torch.zeros(n,
1,
dtype=int,
device=self.__device),
repetition_indices=torch.zeros(n,
dtype=int,
device=self.__device),
)
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 _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 test_spn_sampling(self):
# Define SPN
leaf = distributions.Normal(in_features=2**3,
out_channels=5,
num_repetitions=1)
sum_1 = layers.Sum(in_channels=5,
in_features=2**3,
out_channels=20,
num_repetitions=1)
prd_1 = layers.Product(in_features=2**3,
cardinality=2,
num_repetitions=1)
sum_2 = layers.Sum(in_channels=20,
in_features=2**2,
out_channels=20,
num_repetitions=1)
prd_2 = layers.Product(in_features=2**2,
cardinality=2,
num_repetitions=1)
sum_3 = layers.Sum(in_channels=20,
in_features=2**1,
out_channels=20,
num_repetitions=1)
prd_3 = layers.Product(in_features=2**1,
cardinality=2,
num_repetitions=1)
sum_4 = layers.Sum(in_channels=20,
in_features=2**0,
out_channels=1,
num_repetitions=1)
# Test forward pass
x_test = torch.randn(1, 2**3)
x_test = leaf(x_test)
x_test = sum_1(x_test)
x_test = prd_1(x_test)
x_test = sum_2(x_test)
x_test = prd_2(x_test)
x_test = sum_3(x_test)
x_test = prd_3(x_test)
res = sum_4(x_test)
# Sampling pass
ctx = SamplingContext(n=1000)
sum_4.sample(context=ctx)
prd_3.sample(context=ctx)
sum_3.sample(context=ctx)
prd_2.sample(context=ctx)
sum_2.sample(context=ctx)
prd_1.sample(context=ctx)
sum_1.sample(context=ctx)
samples = leaf.sample(context=ctx)
def sample(self,
n: int = None,
context: SamplingContext = None) -> torch.Tensor:
context = self.prod.sample(context=context)
# Remove padding
if self._pad:
context.parent_indices = context.parent_indices[:, :-self._pad]
samples = self.base_leaf.sample(context=context)
return samples
def test_sum_shape_as_root_node(self):
"""Check that the sum node has the correct sampling shape when used as root."""
n = 5
num_repetitions = 1
for in_channels in [1, 5, 10]:
for in_features in [1, 5, 10]:
sum_layer = layers.Sum(in_channels=in_channels,
out_channels=1,
in_features=in_features,
num_repetitions=num_repetitions)
ctx = SamplingContext(n=n)
ctx = sum_layer.sample(context=ctx)
self.assertTrue(ctx.parent_indices.shape[0] == n)
self.assertTrue(ctx.parent_indices.shape[1] == in_features)
def test_prod_as_intermediate_node(self):
# Product layer values
in_features = 10
num_samples = 5
num_repetitions = 5
for cardinality in range(2, in_features):
prod_layer = layers.Product(in_features=in_features,
cardinality=cardinality,
num_repetitions=num_repetitions)
# Example parent indexes
parent_indices = torch.randint(high=5,
size=(num_samples, in_features))
# Create expected indexes: each index is repeated #cardinality times
pad = (cardinality - in_features % cardinality) % cardinality
expected_sample_indices = []
for j in range(num_samples):
sample_i_indices = []
for i in parent_indices[j, :]:
sample_i_indices += [i] * cardinality
# Remove padding
if pad > 0:
sample_i_indices = sample_i_indices[:-pad]
# Add current sample
expected_sample_indices.append(sample_i_indices)
# As tensor
expected_sample_indices = torch.tensor(expected_sample_indices)
# Sample
ctx = SamplingContext(n=num_samples, parent_indices=parent_indices)
prod_layer.sample(context=ctx)
self.assertTrue(
(expected_sample_indices == ctx.parent_indices).all())
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
def test_spn_mpe(self):
# Define SPN
leaf = distributions.Normal(in_features=2**3,
out_channels=5,
num_repetitions=1)
sum_1 = layers.Sum(in_channels=5,
in_features=2**3,
out_channels=20,
num_repetitions=1)
prd_1 = layers.Product(in_features=2**3,
cardinality=2,
num_repetitions=1)
sum_2 = layers.Sum(in_channels=20,
in_features=2**2,
out_channels=20,
num_repetitions=1)
prd_2 = layers.Product(in_features=2**2,
cardinality=2,
num_repetitions=1)
sum_3 = layers.Sum(in_channels=20,
in_features=2**1,
out_channels=20,
num_repetitions=1)
prd_3 = layers.Product(in_features=2**1,
cardinality=2,
num_repetitions=1)
sum_4 = layers.Sum(in_channels=20,
in_features=2**0,
out_channels=1,
num_repetitions=1)
sum_1._enable_input_cache()
sum_2._enable_input_cache()
sum_3._enable_input_cache()
sum_4._enable_input_cache()
# Test forward pass
x_test = torch.randn(1, 2**3)
x_test = leaf(x_test)
x_test = sum_1(x_test)
x_test = prd_1(x_test)
x_test = sum_2(x_test)
x_test = prd_2(x_test)
x_test = sum_3(x_test)
x_test = prd_3(x_test)
res = sum_4(x_test)
ctx = SamplingContext(n=x_test.shape[0], is_mpe=True)
sum_4.sample(context=ctx)
prd_3.sample(context=ctx)
sum_3.sample(context=ctx)
prd_2.sample(context=ctx)
sum_2.sample(context=ctx)
prd_1.sample(context=ctx)
sum_1.sample(context=ctx)
# Should be the same
mpe_1 = leaf.sample(context=ctx)
mpe_2 = leaf.sample(context=ctx)
mpe_3 = leaf.sample(context=ctx)
self.assertTrue(((mpe_1 - mpe_2).abs() < 1e-6).all())
self.assertTrue(((mpe_2 - mpe_3).abs() < 1e-6).all())
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
# Create sampling context if this is a root layer
if context is None:
assert oc == 1 and r == 1, "Cannot start sampling from non-root layer."
context = SamplingContext(n=n,
parent_indices=None,
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)
n = context.n
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_sampling_input_cache_enabled and self._sampling_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._sampling_input_cache[
i, :, :, context.repetition_indices[i]]
# Create categorical distribution and use weights as logits
dist = torch.distributions.Categorical(logits=log_weights)
indices = dist.sample()
# Update parent indices
context.parent_indices = indices
return context