Closed Pol-Zeimet closed 3 years ago
In principle when you build GraphsTuples on your own, you are in charge of adding the offsets yourself.
We are working on a beta version of a method that batches tf.Dataset
with GraphsTuples correctly by adding the offsets, following similar logic to
utils_tf.concat(..., axis=0)`(link), but it is not yet ready to be open sourced.
In the meantime, you may want to use the hidden method utils_tf._compute_stacked_offsets
(link, and add the output of that to the senders and receivers before putting them in the graphs tuple.
Actually, since the number of edges and nodes seems to be known and fixed statically, you can just add:
num_previous_accumulated_nodes_per_graph = tf.range(0, batch_size) * num_nodes
offsets = tf.tile(num_previous_accumulated_nodes_per_graph[:, None], [1, num_edges])
offsets = tf.reshape(offsets, [-1])
to the senders and receivers. (where num_nodes and num_edges are the number of nodes and number of edges per graph.
The explains the behaviour of my code :) Than you for also posting a possible solution for the offset 🥇 I might have to adjust it to account for the -1 padding when applying it to my senders/receivers. But you really saved me a lot of time!
Quick edit for final code. Now working as expected.
num_previous_accumulated_nodes_per_graph = tf.range(0, batch_size) * num_nodes
offsets = tf.tile(num_previous_accumulated_nodes_per_graph[:, None], [1, num_edges])
offsets = tf.reshape(offsets, [-1])
senders = tf.reshape(senders, [batch_size * num_edges])
receivers = tf.reshape(receivers, [batch_size * num_edges])
senders_offset = tf.where(senders != -1, senders + offsets, senders)
receivers_offset = tf.where(receivers != -1, receivers + offsets, receivers)
combined_graphs_tuple = graphs.GraphsTuple(n_node =tf.fill([batch_size], num_nodes),
n_edge=tf.fill([batch_size], num_edges),
nodes=tf.reshape(nodes, [batch_size * num_nodes, self.node_shape]),
edges=tf.reshape(edges, [batch_size * num_edges, self.edges_shape]),
senders=senders_offset,
receivers=receivers_offset,
globals=None,
)
I am currently working on a project which shares some similarities with the code posted in this issue.
The Idea is to build a
GraphsTuple
from batches of nodes, edges, senders an receivers. Connected nodes an edges will be concatenated and, via linear layers, turned into new node embeddings. This all works as it should with a batch size of 1.As soon as more graphs are involved, the
GraphsTuple
does not behave as I expected. Instead of aggregating nodes and edges according to the individual graphs, the aggregation interprets senders and receivers as if they only affect the first graph in the batch. i.e. if senders and receivers ofgraphs[1]
contain "0", it will addressnodes[0]
of the combinedGraphsTuple
and thereforenode[0]
ofgraphs[0]
.This absolutely makes sense, except I expected the senders and receivers to be automatically offset in this context due to the provided
n_nodes
andn_edges
. i.e. senders of graph[1] becoming offset viasenders = senders + n_nodes[0]
Or am I missing something obvious here?
Here is my code.
Maybe I am building my Tuple in the wrong way?
If you want to test it, you can look at and compare the output like this:
Running it with my data, the wirst output of new_nodes,new_edges,new_senders,new_receivers looks just as it should The second one however does not. The tensor for new_nodes_batched for example contains one array that looks just fine, the second one however is filled with 0
Sorry for the wall of text and thank you in advance :)