Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
dgl
Commits
107b4347
Unverified
Commit
107b4347
authored
Jan 05, 2024
by
Rhett Ying
Committed by
GitHub
Jan 05, 2024
Browse files
[example] update docstring (#6902)
parent
566910d8
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
8 additions
and
2 deletions
+8
-2
examples/sampling/graphbolt/rgcn/hetero_rgcn.py
examples/sampling/graphbolt/rgcn/hetero_rgcn.py
+8
-2
No files found.
examples/sampling/graphbolt/rgcn/hetero_rgcn.py
View file @
107b4347
...
...
@@ -136,7 +136,11 @@ def create_dataloader(
# Move the mini-batch to the appropriate device.
# `device`:
# The device to move the mini-batch to.
# [TODO] Moving `MiniBatch` to GPU is not supported yet.
# [Rui] Usually, we move the mini-batch to target device in the datapipe.
# However, in this example, we leaves the mini-batch on CPU and move it to
# GPU after blocks are created. This is because this example is busy on
# GPU due to embedding layer. And block creation on CPU could be overlapped
# with optimization operation on GPU and it results in better performance.
device
=
torch
.
device
(
"cpu"
)
datapipe
=
datapipe
.
copy_to
(
device
)
...
...
@@ -443,6 +447,7 @@ def evaluate(
y_true
=
list
()
for
data
in
tqdm
(
data_loader
,
desc
=
"Inference"
):
# Convert MiniBatch to DGL Blocks and move them to the target device.
blocks
=
[
block
.
to
(
device
)
for
block
in
data
.
blocks
]
node_features
=
extract_node_features
(
name
,
blocks
[
0
],
data
,
node_embed
,
device
...
...
@@ -508,7 +513,8 @@ def train(
total_loss
=
0
for
data
in
tqdm
(
data_loader
,
desc
=
f
"Training~Epoch
{
epoch
+
1
:
02
d
}
"
):
# Convert MiniBatch to DGL Blocks.
# Convert MiniBatch to DGL Blocks and move them to the target
# device.
blocks
=
[
block
.
to
(
device
)
for
block
in
data
.
blocks
]
# Fetch the number of seed nodes in the batch.
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment