pages/19_Graphs3.py

import streamlit as st
import tensorflow as tf
import tensorflow_gnn as tfgnn
from tensorflow_gnn import runner
from tensorflow_gnn.experimental import sampler
from tensorflow_gnn.models import mt_albis
import functools
import os
from typing import Mapping

# Set environment variable for legacy Keras
os.environ['TF_USE_LEGACY_KERAS'] = '1'

# Set Streamlit title
st.title("Solving OGBN-MAG end-to-end with TF-GNN")

st.write("Setting up the environment...")
tf.get_logger().setLevel('ERROR')
st.write(f"Running TF-GNN {tfgnn.__version__} under TensorFlow {tf.__version__}.")

NUM_TRAINING_SAMPLES = 629571
NUM_VALIDATION_SAMPLES = 64879

GRAPH_TENSOR_FILE = 'gs://download.tensorflow.org/data/ogbn-mag/sampled/v2/graph_tensor.example.pb'
SCHEMA_FILE = 'gs://download.tensorflow.org/data/ogbn-mag/sampled/v2/graph_schema.pbtxt'

# Load the graph schema and graph tensor
st.write("Loading graph schema and tensor...")
graph_schema = tfgnn.read_schema(SCHEMA_FILE)
serialized_ogbn_mag_graph_tensor_string = tf.io.read_file(GRAPH_TENSOR_FILE)

full_ogbn_mag_graph_tensor = tfgnn.parse_single_example(
    tfgnn.create_graph_spec_from_schema_pb(graph_schema, indices_dtype=tf.int64),
    serialized_ogbn_mag_graph_tensor_string)

st.write("Graph tensor loaded successfully.")

# Define sampling sizes
train_sampling_sizes = {
    "cites": 8,
    "rev_writes": 8,
    "writes": 8,
    "affiliated_with": 8,
    "has_topic": 8,
}
validation_sample_sizes = train_sampling_sizes.copy()

# Create sampling model
def create_sampling_model(full_graph_tensor: tfgnn.GraphTensor, sizes: Mapping[str, int]) -> tf.keras.Model:
    def edge_sampler(sampling_op: tfgnn.sampler.SamplingOp):
        edge_set_name = sampling_op.edge_set_name
        sample_size = sizes[edge_set_name]
        return sampler.InMemUniformEdgesSampler.from_graph_tensor(
            full_graph_tensor, edge_set_name, sample_size=sample_size
        )

    def get_features(node_set_name: tfgnn.NodeSetName):
        return sampler.InMemIndexToFeaturesAccessor.from_graph_tensor(
            full_graph_tensor, node_set_name
        )

    # Spell out the sampling procedure in python
    sampling_spec_builder = tfgnn.sampler.SamplingSpecBuilder(graph_schema)
    seed = sampling_spec_builder.seed("paper")
    papers_cited_from_seed = seed.sample(sizes["cites"], "cites")
    authors_of_papers = papers_cited_from_seed.join([seed]).sample(sizes["rev_writes"], "rev_writes")
    papers_by_authors = authors_of_papers.sample(sizes["writes"], "writes")
    institutions = authors_of_papers.sample(sizes["affiliated_with"], "affiliated_with")
    fields_of_study = seed.join([papers_cited_from_seed, papers_by_authors]).sample(sizes["has_topic"], "has_topic")
    sampling_spec = sampling_spec_builder.build()

    model = sampler.create_sampling_model_from_spec(
        graph_schema, sampling_spec, edge_sampler, get_features,
        seed_node_dtype=tf.int64)

    return model

# Create the sampling model
st.write("Creating sampling model...")
sampling_model = create_sampling_model(full_ogbn_mag_graph_tensor, train_sampling_sizes)

st.write("Sampling model created successfully.")

# Define seed dataset function
def seed_dataset(years: tf.Tensor, split_name: str) -> tf.data.Dataset:
    """Seed dataset as indices of papers within split years."""
    if split_name == "train":
        mask = years <= 2017  # 629,571 examples
    elif split_name == "validation":
        mask = years == 2018  # 64,879 examples
    elif split_name == "test":
        mask = years == 2019  # 41,939 examples
    else:
        raise ValueError(f"Unknown split_name: '{split_name}'")
    seed_indices = tf.squeeze(tf.where(mask), axis=-1)
    return tf.data.Dataset.from_tensor_slices(seed_indices)

# Define SubgraphDatasetProvider
class SubgraphDatasetProvider(runner.DatasetProvider):
    """Dataset Provider based on Sampler V2."""

    def __init__(self, full_graph_tensor: tfgnn.GraphTensor, sizes: Mapping[str, int], split_name: str):
        super().__init__()
        self._years = tf.squeeze(full_graph_tensor.node_sets["paper"]["year"], axis=-1)
        self._sampling_model = create_sampling_model(full_graph_tensor, sizes)
        self._split_name = split_name
        self.input_graph_spec = self._sampling_model.output.spec

    def get_dataset(self, context: tf.distribute.InputContext) -> tf.data.Dataset:
        """Creates TF dataset."""
        self._seed_dataset = seed_dataset(self._years, self._split_name)
        ds = self._seed_dataset.shard(
            num_shards=context.num_input_pipelines, index=context.input_pipeline_id)
        if self._split_name == "train":
            ds = ds.shuffle(NUM_TRAINING_SAMPLES).repeat()
        ds = ds.batch(128)
        ds = ds.map(
            functools.partial(self.sample),
            num_parallel_calls=tf.data.AUTOTUNE,
            deterministic=False,
        )
        return ds.unbatch().prefetch(tf.data.AUTOTUNE)

    def sample(self, seeds: tf.Tensor) -> tfgnn.GraphTensor:
        seeds = tf.cast(seeds, tf.int64)
        batch_size = tf.size(seeds)
        seeds_ragged = tf.RaggedTensor.from_row_lengths(
            seeds, tf.ones([batch_size], tf.int64),
        )
        return self._sampling_model(seeds_ragged)

# Create dataset providers
st.write("Creating dataset providers...")
train_ds_provider = SubgraphDatasetProvider(full_ogbn_mag_graph_tensor, train_sampling_sizes, "train")
valid_ds_provider = SubgraphDatasetProvider(full_ogbn_mag_graph_tensor, validation_sample_sizes, "validation")
example_input_graph_spec = train_ds_provider.input_graph_spec._unbatch()

st.write("Dataset providers created successfully.")

# Define the model function
node_state_dim = 128
num_graph_updates = 4
message_dim = 128
state_dropout_rate = 0.2
l2_regularization = 1e-5

def set_initial_node_states(node_set: tfgnn.NodeSet, node_set_name: str):
    if node_set_name == "field_of_study":
        return tf.keras.layers.Embedding(50_000, 32)(node_set["hashed_id"])
    if node_set_name == "institution":
        return tf.keras.layers.Embedding(6_500, 16)(node_set["hashed_id"])
    if node_set_name == "paper":
        return tf.keras.layers.Dense(node_state_dim, activation="relu")(node_set["feat"])
    if node_set_name == "author":
        return node_set["empty_state"]
    raise KeyError(f"Unexpected node_set_name='{node_set_name}'")

def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
    inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)
    graph = tfgnn.keras.layers.MapFeatures(node_sets_fn=set_initial_node_states)(inputs)
    
    for _ in range(num_graph_updates):
        graph = mt_albis.MtAlbisGraphUpdate(
            units=node_state_dim,
            message_dim=message_dim,
            attention_type="none",
            simple_conv_reduce_type="mean|sum",
            normalization_type="layer",
            next_state_type="residual",
            state_dropout_rate=state_dropout_rate,
            l2_regularization=l2_regularization
        )(graph)
    
    paper_state = tfgnn.keras.layers.Readout(node_set_name="paper", feature_name="state")(graph)
    paper_state = tf.keras.layers.Dense(349, activation="softmax")(paper_state)
    return tf.keras.Model(inputs, paper_state)

# Check for TPU/ GPU and set strategy
st.write("Setting up strategy for distributed training...")
if tf.config.list_physical_devices("TPU"):
    st.write("Using TPUStrategy")
    strategy = runner.TPUStrategy("local")
    train_padding = runner.FitOrSkipPadding(example_input_graph_spec, train_ds_provider)
    valid_padding = runner.TightPadding(example_input_graph_spec, valid_ds_provider)
elif tf.config.list_physical_devices("GPU"):
    st.write("Using MirroredStrategy for GPUs")
    strategy = tf.distribute.MirroredStrategy()
    train_padding = None
    valid_padding = None
else:
    st.write("Using default strategy")
    strategy = tf.distribute.get_strategy()
    train_padding = None
    valid_padding = None

st.write(f"Found {strategy.num_replicas_in_sync} replicas in sync")

# Define task
st.write("Defining the task...")
task = runner.NodeMulticlassClassification(
    num_classes=349,
    label_feature_name="paper_venue")

# Set hyperparameters
st.write("Setting hyperparameters...")
global_batch_size = 128
epochs = 10
initial_learning_rate = 0.001

steps_per_epoch = NUM_TRAINING_SAMPLES // global_batch_size
validation_steps = NUM_VALIDATION_SAMPLES // global_batch_size
learning_rate = tf.keras.optimizers.schedules.CosineDecay(
    initial_learning_rate, steps_per_epoch * epochs)
optimizer_fn = functools.partial(tf.keras.optimizers.Adam, learning_rate=learning_rate)

# Define trainer
st.write("Setting up the trainer...")
trainer = runner.KerasTrainer(
    strategy=strategy,
    model_dir="/tmp/gnn_model/",
    callbacks=None,
    steps_per_epoch=steps_per_epoch,
    validation_steps=validation_steps,
    restore_best_weights=False,
    checkpoint_every_n_steps="never",
    summarize_every_n_steps="never",
    backup_and_restore=False,
)

# Run training
st.write("Training the model...")
runner.run(
    task=task,
    model_fn=model_fn,
    trainer=trainer,
    optimizer_fn=optimizer_fn,
    epochs=epochs,
    global_batch_size=global_batch_size,
    train_ds_provider=train_ds_provider,
    valid_ds_provider=valid_ds_provider,
    gtspec=example_input_graph_spec,
)

st.write("Training completed successfully.")