"docs/vscode:/vscode.git/clone" did not exist on "e2b7f745cac29802d596b607ee99058ef8ce17e2"
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 ...@@ -43,6 +43,7 @@ main
""" """
import argparse import argparse
import time import time
from functools import partial
import dgl.graphbolt as gb import dgl.graphbolt as gb
import dgl.nn as dglnn import dgl.nn as dglnn
...@@ -191,7 +192,9 @@ def create_dataloader(args, graph, features, itemset, is_train=True): ...@@ -191,7 +192,9 @@ def create_dataloader(args, graph, features, itemset, is_train=True):
# the negative samples. # the negative samples.
############################################################################ ############################################################################
if is_train and args.exclude_edges: 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]: # [Input]:
......
...@@ -349,10 +349,10 @@ def train( ...@@ -349,10 +349,10 @@ def train(
use_uva=use_uva, use_uva=use_uva,
) )
opt = torch.optim.Adam(model.parameters(), lr=args.lr) opt = torch.optim.Adam(model.parameters(), lr=args.lr)
start_epoch_time = time.time()
for epoch in range(args.epochs): for epoch in range(args.epochs):
model.train() model.train()
total_loss = 0 total_loss = 0
start_epoch_time = time.time()
# A block is a graph consisting of two sets of nodes: the # A block is a graph consisting of two sets of nodes: the
# source nodes and destination nodes. The source and destination # source nodes and destination nodes. The source and destination
# nodes can have multiple node types. All the edges connect from # 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