[GraphBolt] add exclude_seed_edges() for link prediction example (#6765)

a4cacf13 · Rhett Ying · GitHub · d57564ac · a4cacf13
Unverified Commit a4cacf13 authored Dec 18, 2023 by Rhett Ying Committed by GitHub Dec 18, 2023
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 0 deletions

notebooks/stochastic_training/link_prediction.ipynb notebooks/stochastic_training/link_prediction.ipynb +2 -0

No files found.
--- a/notebooks/stochastic_training/link_prediction.ipynb
+++ b/notebooks/stochastic_training/link_prediction.ipynb
@@ -137,9 +137,11 @@
    {
      "cell_type": "code",
      "source": [
+        "from functools import partial\n",
        "datapipe = gb.ItemSampler(train_set, batch_size=256, shuffle=True)\n",
        "datapipe = datapipe.sample_uniform_negative(graph, 5)\n",
        "datapipe = datapipe.sample_neighbor(graph, [5, 5, 5])\n",
+        "datapipe = datapipe.transform(partial(gb.exclude_seed_edges, include_reverse_edges=True))\n",
        "datapipe = datapipe.fetch_feature(feature, node_feature_keys=[\"feat\"])\n",
        "datapipe = datapipe.copy_to(device)\n",
        "train_dataloader = gb.DataLoader(datapipe, num_workers=0)"