Unverified Commit 5a3c9192 authored by peizhou001's avatar peizhou001 Committed by GitHub
Browse files

[Graphbolt[ Add reverse edgs in link prediction example (#6751)


Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-21-218.ap-northeast-1.compute.internal>
parent 90c3c761
......@@ -43,6 +43,7 @@ main
"""
import argparse
import time
from functools import partial
import dgl.graphbolt as gb
import dgl.nn as dglnn
......@@ -191,7 +192,9 @@ def create_dataloader(args, graph, features, itemset, is_train=True):
# the negative samples.
############################################################################
if is_train and args.exclude_edges:
datapipe = datapipe.transform(gb.exclude_seed_edges)
datapipe = datapipe.transform(
partial(gb.exclude_seed_edges, include_reverse_edges=True)
)
############################################################################
# [Input]:
......
......@@ -349,10 +349,10 @@ def train(
use_uva=use_uva,
)
opt = torch.optim.Adam(model.parameters(), lr=args.lr)
start_epoch_time = time.time()
for epoch in range(args.epochs):
model.train()
total_loss = 0
start_epoch_time = time.time()
# A block is a graph consisting of two sets of nodes: the
# source nodes and destination nodes. The source and destination
# nodes can have multiple node types. All the edges connect from
......
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