TF-GNN error when training a graph: ValueError: Could not unbatch scalar (rank=0) GraphPiece

Hi,
I am trying to model an autoencoder of graph data to best reconstruct the graph but Im facing an issue when trying to do the training.

I created the graph tensor as follows:

edge_sources = train_data['from_encoded'].values
edge_targets = train_data['to_encoded'].values
edge_signs = train_data['sign'].values


node_ids = tf.concat([edge_sources, edge_targets], axis=0)
unique_ids, _ = tf.unique(node_ids)

num_nodes = len(unique_ids)

hidden_state_dim = 64

# Create a GraphTensor from edges and node features
graph = tf_gnn.GraphTensor.from_pieces(
    node_sets={
        "nodes": tf_gnn.NodeSet.from_fields(sizes=[num_nodes], features={'id': unique_ids, 'hidden_state': tf.zeros((num_nodes, hidden_state_dim))})
    },
    edge_sets={
        "edges": tf_gnn.EdgeSet.from_fields(
            sizes=[len(edge_sources)],
            adjacency=tf_gnn.Adjacency.from_indices(
                source=("nodes", edge_sources),
                target=("nodes", edge_targets)
            ),
            features={
                "sign": tf.convert_to_tensor(edge_signs, dtype=tf.float32)
            }
        )
    }
)

Then I created the model this way:

graph_tensor_spec = graph.spec

# Define the GCN model with specified hidden layers
gcn_model = gcn.GCNConv(
        units=64,  # Example hidden layer sizes
        activation='relu',
        use_bias=True
    )
    
# Input layer using the graph tensor spec
inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)

# Apply the GCN model to the inputs
graph_setup = gcn_model(inputs,  edge_set_name="edges")
    
# Extract node states
node_states = graph_setup

decoder = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(64, activation='sigmoid')
])

decoded = decoder(node_states)
#decoded = decoder(embeddings)

autoencoder = tf.keras.Model(inputs=inputs, outputs=decoded)

I setup the training as follows:

autoencoder.compile(optimizer='adam', loss=tf.keras.losses.BinaryCrossentropy(),
    metrics=[tf.keras.metrics.AUC()])
autoencoder.fit(
    x=graph,
    y=graph,  # For autoencoders, input = output
    epochs=1   # Number of training epochs
)

But when I do that, I get the following error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-38-c081c4864baa> in <cell line: 1>()
----> 1 autoencoder.fit(
      2     x=graph,
      3     y=graph,  # For autoencoders, input = output
      4     epochs=1   # Number of training epochs
      5 )

1 frames
/usr/local/lib/python3.10/dist-packages/tensorflow_gnn/graph/graph_piece.py in _unbatch(self)
    780     """Extension Types API: Unbatching."""
    781     if self.rank == 0:
--> 782       raise ValueError('Could not unbatch scalar (rank=0) GraphPiece.')
    783 
    784     def unbatch_fn(spec):

ValueError: Could not unbatch scalar (rank=0) GraphPiece.

I am assuming that the model expects a batched format even though I only have 1 graph. But im not sure how to fix this.

I don’t see any documentation with regard to this.

I tried to train the graph but got an unbatched error which im not sure how to fix. I also suspect that it could be the way I’m feeding the data to “autoencoder.fit()” , maybe I’m not providing it in the right format but I’m not sure since the graph data is a Graph.Tensor

hi @Sid72

is this any of assignment codes that you are trying locally and making changes?

In case no, your error is stating basically the graph codes you wrote using the edge sets are rank to 0 and hence throwing this error.

As you haven’t given information about why you used x as well as u as graph and what kind graph scale you are creating, cannot help much.

Regards
DP