def stack_batch(graph_batch, n_global_features, n_node_features,
n_edge_features):
globals_features = tf.reshape(graph_batch.globals,
shape=[-1, n_global_features])
nodes_features = tf.reshape(graph_batch.nodes, shape=[-1, n_node_features])
edges_features = tf.reshape(graph_batch.edges, shape=[-1, n_edge_features])
receivers = tf.reshape(graph_batch.receivers, shape=[-1])
senders = tf.reshape(graph_batch.senders, shape=[-1])
n_node = tf.reshape(graph_batch.n_node, shape=[-1])
n_edge = tf.reshape(graph_batch.n_edge, shape=[-1])
mask = tf.repeat(
tf.multiply(tf.range(utils_tf.get_num_graphs(graph_batch)), n_node),
repeats=n_edge,
)
receivers = receivers + mask
senders = senders + mask
graph_batch = gn.graphs.GraphsTuple(
nodes=nodes_features,
edges=edges_features,
globals=globals_features,
receivers=receivers,
senders=senders,
n_node=n_node,
n_edge=n_edge,
)
return graph_batch
def __call__(self, graph):
_validate_graph(graph, (EDGES, ),
additional_message="when aggregating from edges.")
num_graphs = utils_tf.get_num_graphs(graph)
graph_index = tf.range(num_graphs)
indices = utils_tf.repeat(graph_index, graph.n_edge, axis=0)
return self._reducer(graph.edges, indices, num_graphs)
def equal_graphs(graphs_a, graphs_b, verbose=True):
cond = tf.constant([True])
v_a = utils_tf.get_num_graphs(graphs_a)
v_b = utils_tf.get_num_graphs(graphs_b)
cond = tf.math.logical_and(cond, v_a == v_b)
if not cond and verbose:
pprint("num_graphs", cond.numpy())
for field in [
"n_node",
"n_edge",
"globals",
"nodes",
"edges",
"receivers",
"senders",
]:
v_a = getattr(graphs_a, field)
v_b = getattr(graphs_b, field)
check_values = tf.reduce_all(tf.equal(v_a, v_b))
cond = tf.math.logical_and(cond, check_values)
check_shape = tf.reduce_all(tf.equal(v_a.shape, v_b.shape))
cond = tf.math.logical_and(cond, check_shape)
if not cond and verbose:
pprint(
field,
f"values = {check_values.numpy()}",
f"shape = {check_shape.numpy()}",
cond.numpy(),
)
if verbose:
print()
if utils_tf.get_num_graphs(graphs_a) > 1:
for graph_idx in range(utils_tf.get_num_graphs(graphs_a)):
graph_a = utils_tf.get_graph(graphs_a, graph_idx)
graph_b = utils_tf.get_graph(graphs_b, graph_idx)
check = equal_graphs(graph_a, graph_b, verbose=False)
cond = tf.math.logical_and(cond, check)
return cond
def autoregressive_connect_graph_dynamic(
graph,
exclude_self_edges=False,
name="autoregressive_connect_graph_dynamic"):
"""Adds edges to a graph by fully-connecting the nodes.
This method does not require the number of nodes per graph to be constant,
or to be known at graph building time.
Args:
graph: A `graphs.GraphsTuple` with `None` values for the edges, senders and
receivers.
exclude_self_edges (default=False): Excludes self-connected edges.
name: (string, optional) A name for the operation.
Returns:
A `graphs.GraphsTuple` containing `Tensor`s with fully-connected edges.
Raises:
ValueError: if any of the `EDGES`, `RECEIVERS` or `SENDERS` field is not
`None` in `graph`.
"""
utils_tf._validate_edge_fields_are_all_none(graph)
with tf.name_scope(name):
def body(i, senders, receivers, n_edge):
edges = _create_autogressive_edges_from_nodes_dynamic(
graph.n_node[i], exclude_self_edges)
return (i + 1, senders.write(i, edges['senders']),
receivers.write(i, edges['receivers']),
n_edge.write(i, edges['n_edge']))
num_graphs = utils_tf.get_num_graphs(graph)
loop_condition = lambda i, *_: tf.less(i, num_graphs)
initial_loop_vars = [0] + [
tf.TensorArray(dtype=tf.int32, size=num_graphs, infer_shape=False)
for _ in range(3) # senders, receivers, n_edge
]
_, senders_array, receivers_array, n_edge_array = tf.while_loop(
loop_condition, body, initial_loop_vars, back_prop=False)
n_edge = n_edge_array.concat()
offsets = utils_tf._compute_stacked_offsets(graph.n_node, n_edge)
senders = senders_array.concat() + offsets
receivers = receivers_array.concat() + offsets
senders.set_shape(offsets.shape)
receivers.set_shape(offsets.shape)
receivers.set_shape([None])
senders.set_shape([None])
num_graphs = graph.n_node.get_shape().as_list()[0]
n_edge.set_shape([num_graphs])
return graph.replace(senders=senders,
receivers=receivers,
n_edge=n_edge)
def connect_graph_dynamic(graph: GraphsTuple, is_edge_func, name="connect_graph_dynamic"):
"""
Connects a graph using a boolean edge mask to create edges.
Args:
graph: GraphsTuple
is_edge_func: callable(sender: int, receiver: int) -> bool, should broadcast
name:
Returns:
connected GraphsTuple
"""
utils_tf._validate_edge_fields_are_all_none(graph)
with tf.name_scope(name):
def body(i, senders, receivers, n_edge):
edges = _create_functional_connect_edges_dynamic(graph.n_node[i], is_edge_func)
# edges = create_edges_func(graph.n_node[i])
return (i + 1, senders.write(i, edges['senders']),
receivers.write(i, edges['receivers']),
n_edge.write(i, edges['n_edge']))
num_graphs = utils_tf.get_num_graphs(graph)
loop_condition = lambda i, *_: tf.less(i, num_graphs)
initial_loop_vars = [0] + [
tf.TensorArray(dtype=tf.int32, size=num_graphs, infer_shape=False)
for _ in range(3) # senders, receivers, n_edge
]
_, senders_array, receivers_array, n_edge_array = tf.while_loop(loop_condition, body, initial_loop_vars)
n_edge = n_edge_array.concat()
offsets = utils_tf._compute_stacked_offsets(graph.n_node, n_edge)
senders = senders_array.concat() + offsets
receivers = receivers_array.concat() + offsets
senders.set_shape(offsets.shape)
receivers.set_shape(offsets.shape)
receivers.set_shape([None])
senders.set_shape([None])
num_graphs = graph.n_node.get_shape().as_list()[0]
n_edge.set_shape([num_graphs])
return graph.replace(senders=tf.stop_gradient(senders),
receivers=tf.stop_gradient(receivers),
n_edge=tf.stop_gradient(n_edge))
def _parse_single_example(example, options):
"""Parses a single tf.Example proto.
Args:
example: An Example proto.
options: An instance of reader_pb2.Reader.
Returns:
A dictionary indexed by tensor name.
"""
# Initialize `keys_to_features`.
example_fmt = {
'id':
tf.io.FixedLenFeature([], tf.int64, default_value=-1),
# Proposals
'image/n_proposal':
tf.io.FixedLenFeature([], tf.int64, default_value=0),
'image/proposal/bbox/ymin':
tf.io.VarLenFeature(tf.float32),
'image/proposal/bbox/xmin':
tf.io.VarLenFeature(tf.float32),
'image/proposal/bbox/ymax':
tf.io.VarLenFeature(tf.float32),
'image/proposal/bbox/xmax':
tf.io.VarLenFeature(tf.float32),
'image/proposal/feature':
tf.io.VarLenFeature(tf.float32),
# Caption graph.
'caption_graph/caption':
tf.io.VarLenFeature(tf.string),
'caption_graph/n_node':
tf.io.VarLenFeature(tf.int64),
'caption_graph/n_edge':
tf.io.VarLenFeature(tf.int64),
'caption_graph/nodes':
tf.io.VarLenFeature(tf.string),
'caption_graph/edges':
tf.io.VarLenFeature(tf.string),
'caption_graph/senders':
tf.io.VarLenFeature(tf.int64),
'caption_graph/receivers':
tf.io.VarLenFeature(tf.int64),
# Ground-truth.
'scene_graph/n_relation':
tf.io.FixedLenFeature([], tf.int64, default_value=0),
'scene_graph/predicate':
tf.io.VarLenFeature(tf.string),
'scene_graph/subject':
tf.io.VarLenFeature(tf.string),
'scene_graph/subject/bbox/ymin':
tf.io.VarLenFeature(tf.float32),
'scene_graph/subject/bbox/xmin':
tf.io.VarLenFeature(tf.float32),
'scene_graph/subject/bbox/ymax':
tf.io.VarLenFeature(tf.float32),
'scene_graph/subject/bbox/xmax':
tf.io.VarLenFeature(tf.float32),
'scene_graph/object':
tf.io.VarLenFeature(tf.string),
'scene_graph/object/bbox/ymin':
tf.io.VarLenFeature(tf.float32),
'scene_graph/object/bbox/xmin':
tf.io.VarLenFeature(tf.float32),
'scene_graph/object/bbox/ymax':
tf.io.VarLenFeature(tf.float32),
'scene_graph/object/bbox/xmax':
tf.io.VarLenFeature(tf.float32),
}
# Decode proposals.
parsed = tf.parse_single_example(example, example_fmt)
proposals = tfexample_decoder.BoundingBox(
prefix='image/proposal/bbox/').tensors_to_item(parsed)
n_proposal = tf.minimum(parsed['image/n_proposal'], options.max_n_proposal)
proposals = proposals[:options.max_n_proposal, :]
proposal_features = tf.reshape(
tf.sparse_tensor_to_dense(parsed['image/proposal/feature']),
[-1, options.feature_dimensions])[:options.max_n_proposal, :]
# Decode and randomly sample a caption graph.
graphs = GraphsTuple(
globals=None,
n_node=tf.sparse_tensor_to_dense(parsed['caption_graph/n_node'], 0),
n_edge=tf.sparse_tensor_to_dense(parsed['caption_graph/n_edge'], 0),
nodes=tf.sparse_tensor_to_dense(parsed['caption_graph/nodes'], ''),
edges=tf.sparse_tensor_to_dense(parsed['caption_graph/edges'], ''),
senders=tf.sparse_tensor_to_dense(parsed['caption_graph/senders'], 0),
receivers=tf.sparse_tensor_to_dense(parsed['caption_graph/receivers'], 0))
num_graphs = utils_tf.get_num_graphs(graphs)
index = tf.random.uniform([], minval=0, maxval=num_graphs, dtype=tf.int32)
caption = tf.sparse_tensor_to_dense(parsed['caption_graph/caption'])[index]
caption_graph = utils_tf.get_graph(graphs, index)
# Decode the ground-truth.
subject_boxes = tfexample_decoder.BoundingBox(
prefix='scene_graph/subject/bbox/').tensors_to_item(parsed)
object_boxes = tfexample_decoder.BoundingBox(
prefix='scene_graph/object/bbox/').tensors_to_item(parsed)
feature_dict = {
'id':
parsed['id'],
# Proposals.
'image/n_proposal':
n_proposal,
'image/proposal':
proposals,
'image/proposal/feature':
proposal_features,
# Caption graph.
'caption_graph/caption':
caption,
'caption_graph/n_node':
caption_graph.n_node[0],
'caption_graph/n_edge':
caption_graph.n_edge[0],
'caption_graph/nodes':
caption_graph.nodes,
'caption_graph/edges':
caption_graph.edges,
'caption_graph/senders':
caption_graph.senders,
'caption_graph/receivers':
caption_graph.receivers,
# Ground-truth.
'scene_graph/n_relation':
parsed['scene_graph/n_relation'],
'scene_graph/predicate':
tf.sparse_tensor_to_dense(parsed['scene_graph/predicate'], ''),
'scene_graph/subject':
tf.sparse_tensor_to_dense(parsed['scene_graph/subject'], ''),
'scene_graph/subject/box':
subject_boxes,
'scene_graph/object':
tf.sparse_tensor_to_dense(parsed['scene_graph/object'], ''),
'scene_graph/object/box':
object_boxes,
}
for key in feature_dict.keys():
if key != 'id' and feature_dict[key].dtype == tf.int64:
feature_dict[key] = tf.cast(feature_dict[key], tf.int32)
return feature_dict
