In fact, I use dgl.dataloading.DataLoader for edge sampling in the training script.
For the customized negative sampler, it needs to pick up specific edges of the source nodes that meet the condition (e.g., edge feature weights < 1). The code might be like the below:
def get_negative_edges(g, seed_edges):
# get the device of edge id tensor
device = seed_edges.device
# source nodes
src, _ = g.find_edges(seed_edges, etype=etype)
# find all out edges of source nodes
out_edges = [g.out_edges(s, form='eid', etype=etype) for s in src] # list of tensors, length = batch_size
# filter the edges
neg_edges = [out_e[g.edata['weights'][canonical_etype][out_e] < 1] for out_e in out_edges]
# randomly sample one negative edge
neg_list = 
for neg_e in neg_edges:
if neg_e.shape > 0:
tar_ntype = canonical_etype[-1]
neg_list.append(g.nodes(tar_ntype)[torch.randint(0, g.num_nodes(tar_ntype), (1, ))].to(device))
# edge form: eid -> uv
ret = g.find_edges(torch.cat(dst), etype=etype)
It seems that some operations used in the negative sampler could not directly be done with UVA enabled, which means that there might be unnecessary data transport between CPU and GPU.
Now I am trying to rewrite this sampling function. Could you please provide some suggestions for it? Thanks!