def initialize_data_dependent(self, batch_list):
# Batch list needs to consist of tuples: (z, kwargs)
# kwargs contains the adjacency matrix as well
with torch.no_grad():
for batch, kwargs in batch_list:
kwargs["channel_padding_mask"] = create_channel_mask(
kwargs["length"], max_len=batch.shape[1])
for module_index, module_list in enumerate([[self.node_encoding],
self.step1_flows]):
for layer_index, layer in enumerate(module_list):
print("Processing layer %i (module %i)..." %
(layer_index + 1, module_index + 1),
end="\r")
if isinstance(layer, FlowLayer):
batch_list = FlowModel.run_data_init_layer(
batch_list, layer)
elif isinstance(layer, FlowModel):
batch_list = layer.initialize_data_dependent(
batch_list)
else:
print("[!] ERROR: Unknown layer type", layer)
sys.exit(1)
## Initialize main flow
for i in range(len(batch_list)):
z_nodes, kwargs = batch_list[i]
z_adjacency, x_indices, mask_valid = adjacency2pairs(
adjacency=kwargs["adjacency"], length=kwargs["length"])
attr_mask_valid = mask_valid * (z_adjacency != 0).to(
mask_valid.dtype)
z_edges, _, _ = self.edge_attr_encoding(
(z_adjacency - 1).clamp(min=0),
reverse=False,
channel_padding_mask=attr_mask_valid.unsqueeze(dim=-1))
kwargs["original_z_adjacency"] = z_adjacency
kwargs["binary_adjacency"] = (kwargs["adjacency"] > 0).long()
kwargs["original_mask_valid"] = mask_valid
kwargs["mask_valid"] = attr_mask_valid
kwargs["x_indices"] = x_indices
batch_list[i] = ([z_nodes, z_edges], kwargs)
for layer_index, layer in enumerate(self.step2_flows):
batch_list = FlowModel.run_data_init_layer(batch_list, layer)
for i in range(len(batch_list)):
z, kwargs = batch_list[i]
z_nodes, z_edges = z[0], z[1]
no_edge_mask_valid = kwargs["original_mask_valid"] * (
kwargs["original_z_adjacency"] == 0).float()
z_no_edges, _, _ = self.edge_virtual_encoding(
torch.zeros_like(kwargs["original_z_adjacency"]),
reverse=False,
channel_padding_mask=no_edge_mask_valid.unsqueeze(dim=-1))
z_edges = z_edges * (1 - no_edge_mask_valid)[
..., None] + z_no_edges * no_edge_mask_valid[..., None]
kwargs["mask_valid"] = kwargs["original_mask_valid"]
kwargs.pop("binary_adjacency")
batch_list[i] = ([z_nodes, z_edges], kwargs)
for layer_index, layer in enumerate(self.step3_flows):
batch_list = FlowModel.run_data_init_layer(batch_list, layer)
def test_reversibility(self, z_nodes, adjacency, length, **kwargs):
ldj = z_nodes.new_zeros(z_nodes.size(0), dtype=torch.float32)
if length is not None:
kwargs["length"] = length
kwargs["channel_padding_mask"] = create_channel_mask(
length, max_len=z_nodes.size(1))
## Performing encoding of step 1
z_nodes, ldj = self._run_layer(self.node_encoding,
z_nodes,
False,
ldj=ldj,
**kwargs)
z_nodes_embed = z_nodes
ldj_embed = ldj
## Testing step 1 flows
for flow in self.step1_flows:
z_nodes, ldj = self._run_layer(flow,
z_nodes,
reverse=False,
ldj=ldj,
adjacency=adjacency,
**kwargs)
z_nodes_reversed, ldj_reversed = z_nodes, ldj
for flow in reversed(self.step1_flows):
z_nodes_reversed, ldj_reversed = self._run_layer(
flow,
z_nodes_reversed,
reverse=True,
ldj=ldj_reversed,
adjacency=adjacency,
**kwargs)
rev_node = (
(z_nodes_reversed - z_nodes_embed).abs() > 1e-3).sum() == 0 and (
(ldj_reversed - ldj_embed).abs() > 1e-1).sum() == 0
if not rev_node:
print("[#] WARNING: Step 1 - Coupling layers are not precisely reversible. Max diffs:\n" + \
"Nodes: %s\n" % str(torch.max((z_nodes_reversed - z_nodes_embed).abs())) + \
"LDJ: %s" % str(torch.max((ldj_reversed - ldj_embed).abs())))
## Performing encoding of step 2
z_edges_disc, x_indices, mask_valid = adjacency2pairs(
adjacency=adjacency, length=length)
kwargs["mask_valid"] = mask_valid * (z_edges_disc != 0).to(
mask_valid.dtype)
kwargs["x_indices"] = x_indices
binary_adjacency = (adjacency > 0).long()
kwargs_edge_embed = kwargs.copy()
kwargs_edge_embed["channel_padding_mask"] = kwargs[
"mask_valid"].unsqueeze(dim=-1)
z_attr = (z_edges_disc - 1).clamp(min=0)
z_edges, ldj = self._run_layer(self.edge_attr_encoding, z_attr, False,
ldj, **kwargs_edge_embed)
## Testing step 2 flows
z_nodes_orig, z_edges_orig, ldj_orig = z_nodes, z_edges, ldj
for flow in self.step2_flows:
z_nodes, z_edges, ldj = self._run_node_edge_layer(
flow,
z_nodes,
z_edges,
False,
ldj,
binary_adjacency=binary_adjacency,
**kwargs)
z_nodes_rev, z_edges_rev, ldj_rev = z_nodes, z_edges, ldj
for flow in reversed(self.step2_flows):
z_nodes_rev, z_edges_rev, ldj_rev = self._run_node_edge_layer(
flow,
z_nodes_rev,
z_edges_rev,
True,
ldj_rev,
binary_adjacency=binary_adjacency,
**kwargs)
rev_edge_attr = ((z_nodes_rev - z_nodes_orig).abs() > 1e-3).sum() == 0 and \
((z_edges_rev - z_edges_orig).abs() > 1e-3).sum() == 0 and \
((ldj_rev - ldj_orig).abs() > 1e-1).sum() == 0
if not rev_edge_attr:
print("[#] WARNING: Step 2 - Coupling layers are not precisely reversible. Max diffs:\n" + \
"Nodes: %s\n" % str(torch.max((z_nodes_rev - z_nodes_orig).abs())) + \
"Edges: %s\n" % str(torch.max((z_edges_rev - z_edges_orig).abs())) + \
"LDJ: %s" % str(torch.max((ldj_rev - ldj_orig).abs())))
## Performing encoding of step 3
kwargs["mask_valid"] = mask_valid
virtual_edge_mask = mask_valid * (z_edges_disc == 0).float()
kwargs_no_edge_embed = kwargs.copy()
kwargs_no_edge_embed[
"channel_padding_mask"] = virtual_edge_mask.unsqueeze(dim=-1)
virt_edges = z_edges.new_zeros(z_edges.shape[:-1], dtype=torch.long)
z_virtual_edges, ldj = self._run_layer(self.edge_virtual_encoding,
virt_edges, False, ldj,
**kwargs_no_edge_embed)
z_edges = torch.where(
virtual_edge_mask.unsqueeze(dim=-1) == 1, z_virtual_edges, z_edges)
## Testing step 3 flows
z_nodes_orig, z_edges_orig, ldj_orig = z_nodes, z_edges, ldj
for flow in self.step3_flows:
z_nodes, z_edges, ldj = self._run_node_edge_layer(
flow, z_nodes, z_edges, False, ldj, **kwargs)
z_nodes_rev, z_edges_rev, ldj_rev = z_nodes, z_edges, ldj
for flow in reversed(self.step3_flows):
z_nodes_rev, z_edges_rev, ldj_rev = self._run_node_edge_layer(
flow, z_nodes_rev, z_edges_rev, True, ldj_rev, **kwargs)
rev_edge_virt = ((z_nodes_rev - z_nodes_orig).abs() > 1e-3).sum() == 0 and \
((z_edges_rev - z_edges_orig).abs() > 1e-3).sum() == 0 and \
((ldj_rev - ldj_orig).abs() > 1e-1).sum() == 0
if not rev_edge_virt:
print("[#] WARNING: Step 3 - Coupling layers are not precisely reversible. Max diffs:\n" + \
"Nodes: %s\n" % str(torch.max((z_nodes_rev - z_nodes_orig).abs())) + \
"Edges: %s\n" % str(torch.max((z_edges_rev - z_edges_orig).abs())) + \
"LDJ: %s" % str(torch.max((ldj_rev - ldj_orig).abs())))
if rev_node and rev_edge_attr and rev_edge_virt:
print("Reversibility test succeeded!")
else:
print(
"Reversibility test finished with warnings. Non-reversibility can be due to limited precision in mixture coupling layers"
)
def forward(self,
z,
adjacency=None,
ldj=None,
reverse=False,
get_ldj_per_layer=False,
length=None,
sample_temp=1.0,
**kwargs):
z_nodes = z # Renaming as argument "z" is usually used for the flows, but here it represents the discrete node types
if ldj is None:
ldj = z_nodes.new_zeros(z_nodes.size(0), dtype=torch.float32)
if length is not None:
kwargs["length"] = length
kwargs["channel_padding_mask"] = create_channel_mask(
length, max_len=z_nodes.size(1))
ldj_per_layer = []
if not reverse:
## Step 1 => Encode nodes in latent space and apply RGCN flows
z_nodes, ldj = self._step1_forward(z_nodes, adjacency, ldj, False,
ldj_per_layer, **kwargs)
## Edges are represented as a list (1D tensor), not as a matrix, because we do not want to consider each edge twice
# X_indices is a tuple, where each element has the size of the edge tensor and states the node indices of the corresponding edge
# Mask_valid is a tensor of the same size, but contains 0 for those edges that are not "valid".
# This is when edges are padding elements for graphs of different sizes
# here z_edges_disc is the discrete type of edge type; its shape is [batch_size, (37+1)*37/2], mask_valid is to show the valid edges which should not contain any virtual node
z_edges_disc, x_indices, mask_valid = adjacency2pairs(
adjacency=adjacency, length=length)
kwargs["mask_valid"] = mask_valid * (z_edges_disc != 0).to(
mask_valid.dtype
) # rule out the non-exisitng edges between valid nodes
kwargs["x_indices"] = x_indices
binary_adjacency = (adjacency > 0).long(
) # attention that adjacency matrix itself is type-specific matrix while binary matrix is proximity matrix
## Step 2 => Encode edge attributes in latent space and apply first EdgeGNN flows
z_nodes, z_edges, ldj = self._step2_forward(
z_nodes,
z_edges_disc,
ldj,
False,
ldj_per_layer,
binary_adjacency=binary_adjacency,
**kwargs)
## Step 3 => Encode virtual edges in latent space and apply final EdgeGNN flows
kwargs["mask_valid"] = mask_valid
virtual_edge_mask = mask_valid * (z_edges_disc == 0).float()
z_nodes, z_edges, ldj = self._step3_forward(
z_nodes, z_edges, ldj, False, ldj_per_layer, virtual_edge_mask,
**kwargs)
## Add log probability of adjacency matrix to ldj. Only nodes are considered in the task object
adjacency_log_prob = (self.prior_distribution.log_prob(z_edges) *
mask_valid.unsqueeze(dim=-1)).sum(dim=[1, 2])
ldj = ldj + adjacency_log_prob
ldj_per_layer.append({"adjacency_log_prob": adjacency_log_prob})
else:
z_nodes = z
batch_size, num_nodes = z_nodes.size(0), z_nodes.size(1)
## Sample latent variables for adjacency matrix
mask_valid, x_indices = get_adjacency_indices(num_nodes=num_nodes,
length=length)
kwargs["mask_valid"] = mask_valid
kwargs["x_indices"] = x_indices
z_edges = self.prior_distribution.sample(
shape=(batch_size, mask_valid.size(1),
self.encoding_dim_edges),
temp=sample_temp).to(z.device)
## Reverse step 3 => decode virtual edges
z_nodes, z_edges, ldj, mask_valid = self._step3_forward(
z_nodes, z_edges, ldj, True, ldj_per_layer, **kwargs)
binary_adjacency = pairs2adjacency(num_nodes=num_nodes,
pairs=mask_valid,
length=length,
x_indices=x_indices)
## Reverse step 2 => decode edge attributes
kwargs["mask_valid"] = mask_valid
z_nodes, z_edges, ldj = self._step2_forward(
z_nodes,
z_edges,
ldj,
True,
ldj_per_layer,
binary_adjacency=binary_adjacency,
**kwargs)
adjacency = pairs2adjacency(num_nodes=num_nodes,
pairs=z_edges,
length=length,
x_indices=x_indices)
## Reverse step 1 => decode node types
z_nodes, ldj = self._step1_forward(z_nodes,
adjacency,
ldj,
reverse=True,
ldj_per_layer=ldj_per_layer,
**kwargs)
z_nodes = (z_nodes, adjacency)
if get_ldj_per_layer:
return z_nodes, ldj, ldj_per_layer
else:
return z_nodes, ldj