Skip to content

GitLab

Unverified Commit a4cacf13 authored by Rhett Ying's avatar Rhett Ying Committed by GitHub
Browse files

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

parent d57564ac
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
notebooks/stochastic_training/link_prediction.ipynb
View file @ a4cacf13
...@@ -137,9 +137,11 @@ ...@@ -137,9 +137,11 @@
{ {
"cell_type": "code", "cell_type": "code",
"source": [ "source": [
"from functools import partial\n",
"datapipe = gb.ItemSampler(train_set, batch_size=256, shuffle=True)\n", "datapipe = gb.ItemSampler(train_set, batch_size=256, shuffle=True)\n",
"datapipe = datapipe.sample_uniform_negative(graph, 5)\n", "datapipe = datapipe.sample_uniform_negative(graph, 5)\n",
"datapipe = datapipe.sample_neighbor(graph, [5, 5, 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.fetch_feature(feature, node_feature_keys=[\"feat\"])\n",
"datapipe = datapipe.copy_to(device)\n", "datapipe = datapipe.copy_to(device)\n",
"train_dataloader = gb.DataLoader(datapipe, num_workers=0)" "train_dataloader = gb.DataLoader(datapipe, num_workers=0)"
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment