Updated configuration of OGB-MAG distributed example.

PiperOrigin-RevId: 590871689
tensorflow · Dec 14, 2023 · ab78ca3 · ab78ca3
1 parent 84fd231
commit ab78ca3
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 12 deletions.
diff --git a/examples/sampler/mag/graph_schema.pbtxt b/examples/sampler/mag/graph_schema.pbtxt
@@ -50,7 +50,7 @@ node_sets {
       }
     }
     metadata {
-      filename: "nodes-paper.tfrecords@397"
+      filename: "nodes-paper.tfrecords-?????-of-?????"
     }
   }
 }
@@ -60,7 +60,7 @@ edge_sets {
     source: "author"
     target: "institution"
     metadata {
-      filename: "edges-affiliated_with.tfrecords@30"
+      filename: "edges-affiliated_with.tfrecords-?????-of-?????"
     }
   }
 }
@@ -70,7 +70,7 @@ edge_sets {
     source: "paper"
     target: "paper"
     metadata {
-      filename: "edges-cites.tfrecords@120"
+      filename: "edges-cites.tfrecords-?????-of-?????"
     }
   }
 }
@@ -80,7 +80,7 @@ edge_sets {
     source: "paper"
     target: "field_of_study"
     metadata {
-      filename: "edges-has_topic.tfrecords@226"
+      filename: "edges-has_topic.tfrecords-?????-of-?????"
     }
   }
 }
@@ -90,7 +90,7 @@ edge_sets {
     source: "author"
     target: "paper"
     metadata {
-      filename: "edges-writes.tfrecords@172"
+      filename: "edges-writes.tfrecords-?????-of-?????"
     }
   }
 }
@@ -100,7 +100,7 @@ edge_sets {
     source: "paper"
     target: "author"
     metadata {
-      filename: "edges-writes.tfrecords@172"
+      filename: "edges-writes.tfrecords-?????-of-?????"
       extra {
         key: "edge_type"
         value: "reversed"

diff --git a/examples/sampler/mag/run_mag.sh b/examples/sampler/mag/run_mag.sh
@@ -14,6 +14,18 @@
 # limitations under the License.
 # ==============================================================================
 
+# Samples subgraphs for https://ogb.stanford.edu/docs/nodeprop/#ogbn-mag.
+#
+# Consider `tfgnn_convert_ogb_dataset` to convert OGB data to sharded
+# `TFRecord`s on your local machine and then copy to Google Cloud bucket
+# ${DATA_PATH}/graph, e.g. using `gcloud cp`.
+#
+# Input:
+#   nodes: ${DATA_PATH}/graph/nodes-*.tfrecords-?????-of-?????
+#   edges: ${DATA_PATH}/graph/edges-*.tfrecords-?????-of-?????
+#
+# Result:
+#   Sampled subgraphs: ${DATA_PATH}/outputs/examples.tfrecords-?????-of-?????
 NUM_WORKER_THREADS=2
 MACHINE_TYPE="n1-highmem-2"
 
@@ -24,18 +36,20 @@ MIN_NUM_WORKERS=30
 TIMESTAMP="$(date +"%Y-%m-%d-%H-%M-%S")"
 
 GOOGLE_CLOUD_PROJECT="<Your Google Cloud Project>"
-EXAMPLE_ARTIFACT_DIRECTORY="gs://${GOOGLE_CLOUD_PROJECT}/sampler/${DATASET}/${TIMESTAMP}"
+DATA_PATH="gs://<Bucket with prepared OGB MAG data>"
+
 
-python3 -m tensorflow_gnn.experimental.sampler.beam.sampler \
+tfgnn_sampler \
   --project ${GOOGLE_CLOUD_PROJECT} \
   --region "us-east1" \
   --save_main_session \
   --setup_file "./setup.py" \
-  --graph_schema "gs://<path-to-data-dir>/graph_schema.pbtxt" \
-  --sampling_spec "gs://<path-to-data-dir>/sampling_spec.pbtxt" \
-  --output_samples "${EXAMPLE_ARTIFACT_DIRECTORY}/outputs/examples.tfrecord" \
+  --data_path "${DATA_PATH}/graph" \
+  --graph_schema "graph_schema.pbtxt" \
+  --sampling_spec "sampling_spec.pbtxt" \
+  --output_samples "${DATA_PATH}/outputs/examples.tfrecords@100" \
   --runner DataflowRunner \
-  --temp_location "${EXAMPLE_ARTIFACT_DIRECTORY}/tmp" \
+  --temp_location "${DATA_PATH}/tmp" \
   --machine_type ${MACHINE_TYPE} \
   --experiments "min_num_workers=${MIN_NUM_WORKERS}" \
   --max_num_workers ${MAX_NUM_WORKERS} \