Examples graphs about gnn HOT 18 CLOSED

from typing import Iterator

import networkx as nx
import numpy as np
import tensorflow_gnn as tfgnn
import tensorflow as tf
from ogb.nodeproppred import NodePropPredDataset


def ogb_as_networkx_graph(name: str) -> nx.Graph:
    dataset = NodePropPredDataset(name)
    ogb_graph, labels = dataset[0]
    num_nodes = ogb_graph["num_nodes"]
    ogb_features = ogb_graph["node_feat"]
    ogb_edgelist = ogb_graph["edge_index"]
    ogb_node_indices = np.arange(num_nodes)

    graph = nx.from_edgelist(ogb_edgelist.T)
    data = zip(ogb_node_indices, ogb_features, labels)
    features = {
        node: {"features": features, "label": label}
        for node, features, label in data
    }
    nx.set_node_attributes(graph, values=features)
    return graph


def generate_graph_samples(
    graph: nx.Graph, neighbours: int = 2
) -> Iterator[tfgnn.GraphTensor]:
    for node in graph.nodes:
        subgraph = nx.ego_graph(graph, node, radius=neighbours)

        num_edges = subgraph.number_of_edges()
        adjacency = np.asarray(subgraph.edges)
        edges = tfgnn.EdgeSet.from_fields(
            sizes=[num_edges],
            adjacency=tfgnn.Adjacency.from_indices(
                source=("node", adjacency[:, 0]),
                target=("node", adjacency[:, 1]),
            ),
        )

        data = (
            (data["features"], data["label"])
            for _, data in subgraph.nodes(data=True)
        )
        features, labels = zip(*data)
        num_nodes = subgraph.number_of_nodes()
        nodes = tfgnn.NodeSet.from_fields(
            features={
                "hidden_state": np.asarray(features),
                "labels": np.asarray(labels),
            },
            sizes=[num_nodes],
        )

        graph_tensor = tfgnn.GraphTensor.from_pieces(
            edge_sets={"edges": edges}, node_sets={"nodes": nodes}
        )

        yield graph_tensor


graph = ogb_as_networkx_graph("ogbn-arxiv")
graph_samples = generate_graph_samples(graph, neighbours=2)
graph_tensor = next(graph_samples)
type_spec = graph_tensor.spec

inputs = tf.keras.layers.Input(type_spec=type_spec)
gnn = tfgnn.keras.ConvGNNBuilder(
    lambda edge: tfgnn.keras.layers.SimpleConvolution(
        tf.keras.layers.Dense(32)
    ),
    lambda node: tfgnn.keras.layers.NextStateFromConcat(
        tf.keras.layers.Dense(32)
    ),
)
update = gnn.Convolve()(inputs)
hidden = tfgnn.keras.layers.Readout(node_set_name="nodes")(update)
output = tf.keras.layers.Dense(1, activation="sigmoid")(hidden)
model = tf.keras.Model(inputs, output)

y = model(graph_tensor)
print(y)

import tensorflow as tf
import tensorflow_gnn as tfgnn

tf.debugging.disable_traceback_filtering()

schema = tfgnn.read_schema("tensorflow_gnn/testdata/homogeneous/citrus.pbtxt")
spec = tfgnn.create_graph_spec_from_schema_pb(schema)
graph = tfgnn.random_graph_tensor(spec)

gnn = tfgnn.keras.ConvGNNBuilder(
    lambda edge: tfgnn.keras.layers.SimpleConvolution(tf.keras.layers.Dense(32)),
    lambda node: tfgnn.keras.layers.NextStateFromConcat(tf.keras.layers.Dense(32)),
)

inputs = tf.keras.layers.Input(type_spec=spec)
update = gnn.Convolve()(inputs)
output = tfgnn.keras.layers.Readout(node_set_name="fruits")(update)
model = tf.keras.Model(inputs, output)

model(graph)

from typing import Tuple

import tensorflow as tf
import tensorflow_gnn as tfgnn
from tensorflow_gnn.graph.graph_constants import FieldsNest


def build_graph_tensor(
    num_edges: int = 10, num_nodes: int = 10, num_features: int = 5
) -> tfgnn.GraphTensor:
    adjacency = tf.random.stateless_binomial(
        shape=[num_edges, 2], counts=[0, 1], probs=[0.5, 0.5], seed=[0, 0]
    )
    edges = tfgnn.EdgeSet.from_fields(
        features={"edge_weight": tf.random.normal(shape=[num_features])},
        sizes=[num_edges],
        adjacency=tfgnn.HyperAdjacency.from_indices(
            indices={
                tfgnn.SOURCE: ("node", adjacency[:, 0]),
                tfgnn.TARGET: ("node", adjacency[:, 1]),
            }
        ),
    )

    nodes = tfgnn.NodeSet.from_fields(
        features={"hidden_state": tf.random.normal(shape=[num_nodes, num_features])},
        sizes=[num_nodes],
    )

    return tfgnn.GraphTensor.from_pieces(
        edge_sets={"edge": edges}, node_sets={"node": nodes}
    )


class MultiplyNodeEdge(tf.keras.layers.Layer):
    def __init__(self, edge_feature: str, node_feature: str = tfgnn.SOURCE) -> None:
        super().__init__()
        self.node_feature = node_feature
        self.edge_feature = edge_feature

    def call(self, inputs: Tuple[FieldsNest, FieldsNest, FieldsNest]) -> tf.Tensor:
        edge_inputs, node_inputs, _ = inputs
        return tf.multiply(
            edge_inputs[self.edge_feature], node_inputs[self.node_feature]
        )


graph_tensor = build_graph_tensor()
spec = graph_tensor.spec

input = tf.keras.layers.Input(type_spec=spec)
update = tfgnn.keras.layers.GraphUpdate(
    edge_sets={
        "edge": tfgnn.keras.layers.EdgeSetUpdate(
            next_state=MultiplyNodeEdge(edge_feature="edge_weight"),
            edge_input_feature=["edge_weight"],
        )
    },
    node_sets={
        "node": tfgnn.keras.layers.NodeSetUpdate(
            edge_set_inputs={"edge": tfgnn.keras.layers.Pool(tfgnn.TARGET, "sum")},
            next_state=tfgnn.keras.layers.NextStateFromConcat(
                tf.keras.layers.Dense(16)
            ),
        )
    },
)
graph = update(input)
hidden = tfgnn.keras.layers.Readout(node_set_name="node")(graph)
output = tf.keras.layers.Dense(1, activation="sigmoid")(hidden)
model = tf.keras.Model(input, output)

y = model(graph_tensor)
print(y)