Creating custom Layer for NodeSetUpdate or EdgeSetUpdate

Hi, am trying to create my own custom layer that is passes as "edge_set_inputs" where I get the following error, the purpose of the class is to take a transformation layer(EdgesToNodePoolingLayer) as input and before the input graph is passed through the layer make changes to the edges and edge attributes. I need help in the following 2 places:

Resolving the below error where the transformation layer without any change when attempting to call creating an issue
How to get edges source and target along with edge attributes inside the layer

The error faced is:

TypeError: outer_factory..inner_factory..tf__call() got an unexpected keyword argument 'edge_set_name'

For understanding the code better have attached the graph tensor spec, custom class created, and graph building code below.

Graph Tensor Spec:

graph_tensor_spec = tfgnn.GraphTensorSpec.from_piece_specs(
    context_spec=tfgnn.ContextSpec.from_field_specs(features_spec={
                  'label': tf.TensorSpec(shape=(1,), dtype=tf.int32)
    }),
    node_sets_spec={
        'my_node':
            tfgnn.NodeSetSpec.from_field_specs(
                features_spec={
                    tfgnn.HIDDEN_STATE:
                        tf.TensorSpec((None, 1), tf.float32)
                },
                sizes_spec=tf.TensorSpec((1,), tf.int32))
    },
    edge_sets_spec={
        'my_edge':
            tfgnn.EdgeSetSpec.from_field_specs(
                features_spec={
                    tfgnn.HIDDEN_STATE:
                        tf.TensorSpec((None, 1), tf.float32)
                },
                sizes_spec=tf.TensorSpec((1,), tf.int32),
                adjacency_spec=tfgnn.AdjacencySpec.from_incident_node_sets(
                    'my_node', 'my_node'))
    })

Custom Class Created:

class MyLayer(tf.keras.layers.Layer):
    def __init__(self, transformation, **kwargs):
        super().__init__(**kwargs)
        self._transformation_layer = transformation

    def get_config(self):
        return dict(transformation_layer=self._transformation_layer,
                    **super().get_config())

    def call(self, graph: tfgnn.GraphTensor):
        out = self._transformation_layer(graph)
        return out

graph model code:

model_input_graph_spec, label_spec = train_ds.element_spec

node_dim = 32
edge_dim = 32

def set_initial_node_states(node_set: tfgnn.NodeSet, *, node_set_name: str):
    return tf.keras.layers.Dense(node_dim)(node_set[tfgnn.HIDDEN_STATE])

def set_initial_edge_states(edge_set: tfgnn.EdgeSet, *, edge_set_name: str):
    return tf.keras.layers.Dense(edge_dim)(edge_set[tfgnn.HIDDEN_STATE])

l2_regularization = 6E-6
dropout_rate = 0.2
use_layer_normalization = True

inputs = tf.keras.layers.Input(type_spec=model_input_graph_spec)
graph = tfgnn.keras.layers.MapFeatures(
    node_sets_fn=set_initial_node_states, edge_sets_fn=set_initial_edge_states)(inputs)
graph = graph.merge_batch_to_components()
graph = tfgnn.keras.layers.GraphUpdate(node_sets=
{"my_node": tfgnn.keras.layers.NodeSetUpdate({
    "my_edge": MyLayer(SimpleConv(tf.keras.layers.Dense(32,activation='relu'), "sum", receiver_tag=tfgnn.TARGET))},
    tfgnn.keras.layers.NextStateFromConcat(tf.keras.layers.Dense(128)))},

    edge_sets={"my_edge": tfgnn.keras.layers.EdgeSetUpdate(
        tfgnn.keras.layers.NextStateFromConcat(tf.keras.layers.Dense(128)),
    )}
)(graph)

readout_features = tfgnn.keras.layers.Pool(tfgnn.CONTEXT, "mean", node_set_name="pixel")(graph)

out_layer = tf.keras.layers.Dense(10, activation='softmax')(readout_features)
model = tf.keras.Model(inputs, out_layer)

Any help will be greatly appreciated ! , Please do let me know if any more details is needed for the same.

Hi There,

NodeSetUpdate should receive a convolution layer with the following signature:

new_edge_state = edge_set_update(graph, edge_set_name="...")

So MyLayer should have the edge_set_name in the input args for the call() function, but currently it only has graph as input, hence the failure.

Re: How to get the edge's source and target node in the layer along with edge attributes inside the layer;

After you have the edge_set_name you can access the source and target node id's as well as the edge attributes as below:

graph.edge_sets[].adjacency.source graph.edge_sets[].adjacency.target graph.edge_sets[][]

Also I wonder why you'd need a separate MyLayer layer around SimpleConv, while SimpleConv already handles all the convolutions. SimpleConv will help you configure which sender/receiver and/or edge features you'd prefer to use and transform those for you.

Hope this helps!

Hey, thanks for the response!

By getting edge_set_name the above error is resolved, but we still need clarity on how to access edges and nodes and their attributes and properties. Currently in our ongoing research the edges has to be dropped dynamically and edge attributes have to be modified dynamically for each graph, but since the graphs are merged as components in a batch, we are finding it difficult to get the edges and nodes of each graph separately. We tried to print the graph along with edge source, targets and attributes, which gives output as follows.

GraphTensor(
  context=Context(features={}, sizes=Tensor("model/graph_update/ones_like:0", shape=(None,), dtype=int32), shape=(), indices_dtype=tf.int32),
  node_set_names=['my_node'],
  edge_set_names=['my_edge'])
Tensor("model/input.merge_batch_to_components/Add_1:0", shape=(None,), dtype=int32)
Tensor("model/input.merge_batch_to_components/Add_3:0", shape=(None,), dtype=int32)
Tensor("model/graph_update/edge_set_update/next_state_from_concat_1/dense_2/BiasAdd:0", shape=(None, 128), dtype=float32)

The code used:

class DroNo(tf.keras.layers.Layer):
    def __init__(self,transformation: tf.keras.layers.Layer, **kwargs): 
        super().__init__(**kwargs)
        self._transformation = transformation

    def get_config(self):
        return dict(transformation=self._transformation,
                    **super().get_config())

    def call(self, graph: tfgnn.GraphTensor, edge_set_name):
        print(graph)
        print(graph.edge_sets['my_edge'].adjacency.source)
        print(graph.edge_sets['my_edge'].adjacency.target)
        print(graph.edge_sets['my_edge']['hidden_state'])
        return self._transformation(graph, edge_set_name=edge_set_name)

So kindly help us on how to drop edges and nodes dynamically in each graph and alter edge attributes and node properties.

What's the criteria you want to use for updating edges?

If let's say you want to drop an edge if its edge attribute(hidden_state) is all zeros, you could do something like below;


hidden_tensor_sum = tf.math.reduce_sum(graph.edge_sets['my_edge']['hidden_state'], axis=-1)
zeros_mask = tf.math.not_equal(hidden_tensor_sum, 0.0)

# Note that below will also add '_masked_edge_set' to gt.edge_sets dictionary, that you could delete altogether later
gt = tfgnn.mask_edges(gt, 'my_edge', zeros_mask, '_masked')

You can find the data structures that's used by the GraphTensor to keep the edge-sets's adjacency information from the following classes; Adjacency and HyperAdjacency.

After batching, if you still need to identify a particular graph within a merged batch that's somehow independent of its node-set and edge-set details/features then you could potentially add that graph specific information into the respective graph-pieces(nodes/edges/context etc) that could later be used.

To learn more about the TF-GNN you can find tutorials from various conferences here: https://github.com/tensorflow/gnn/tree/main/examples/tutorials, also particularly more about GraphTensors here in the 2nd part of the talk : https://www.youtube.com/watch?v=JqWROPYeqjA

Hey Neslihans, thanks for the help, in our scenario we are trying to randomly mask/drop 30% of edges or nodes or both for each graph during training. Currently we tried to randomly mask edges (without percentage) to check the working ,using the following code

class DroNo(tf.keras.layers.Layer):
    def __init__(self,transformation: tf.keras.layers.Layer, **kwargs): 
        super().__init__(**kwargs)
        self._transformation = transformation

    def get_config(self):
        return dict(transformation=self._transformation,
                    **super().get_config())

    def call(self, graph: tfgnn.GraphTensor, edge_set_name):
        print(graph)
        print(graph.edge_sets['my_edge'].adjacency.source)
        print(graph.edge_sets['my_edge'].adjacency.target)
        print(graph.edge_sets['my_edge']['hidden_state'])
        hidden_tensor_sum = tf.math.reduce_sum(graph.edge_sets['my_edge']['hidden_state'], axis=-1)
        print(tf.random.uniform(hidden_tensor_sum.shape, minval=0, maxval=2, dtype=tf.int32))
        return self._transformation(graph, edge_set_name=edge_set_name)

and we are getting this error

ValueError: Cannot convert a partially known TensorShape (None,) to a Tensor.

Any help will be highly appreciated to obtain our objective which is to drop 20-30% edges and nodes randomly in each graph during training.

The error message

ValueError: Cannot convert a partially known TensorShape (None,) to a Tensor.

comes from this line of code, I suppose:

print(tf.random.uniform(hidden_tensor_sum.shape, minval=0, maxval=2, dtype=tf.int32))

...and presumably can be fixed by using tf.shape(hidden_tensor_sum) instead of hidden_tensor_sum.shape.

For more information, please visit tensorflow.org/guide/tensor and read the explanations around "You may run across not-fully-specified shapes." I suspect your code is run as part of tf.keras.Model.fit(), which uses tf.function under the hood unless you set model.compile(..., run_eagerly=True).

Those TensorFlow generalities aside, is there a particular defect or feature request you'd like to report for the TF-GNN library? If not, I'd like to close this issue.

For general "how-to" questions, please consult the respective StackOverflow tags [tensorflow] and [tensorflow-gnn].

tensorflow / gnn

Creating custom Layer for NodeSetUpdate or EdgeSetUpdate #552