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