Unable to use 'prob' to sample_neighbors in DistGraph

xcwanAndy · October 29, 2021, 9:39am

Hi all.

I found a problem related to distributed sample neighbors here.

First, I start a DistGraph and insert features ‘prob’ (DistTensor) in it. I want to iteratively modify the ‘prob’ feature during training for adaptive sampling. For DistTensors, they are all stored in g._client rather than in local graph’s edata. However, I found that the source code does not consider the prob in an online mode because it always queries g.local_partition where information is not updated.

github.com

dmlc/dgl/blob/master/python/dgl/distributed/graph_services.py#L379

    
      
                  req_list.append((pid, req))
          
          
# send requests to the remote machine.
          msgseq2pos = None
          if len(req_list) > 0:
              msgseq2pos = send_requests_to_machine(req_list)
          
          
# sample neighbors for the nodes in the local partition.
          res_list = []
          if local_nids is not None:
              src, dst, eids = local_access(g.local_partition, partition_book, local_nids)
              res_list.append(LocalSampledGraph(src, dst, eids))
          
          
# receive responses from remote machines.
          if msgseq2pos is not None:
              results = recv_responses(msgseq2pos)
              res_list.extend(results)
          
          
sampled_graph = merge_graphs(res_list, g.number_of_nodes())
          return sampled_graph

This means that ‘prob’ is not yet available in DistGraph, right?

Is there any suggestions for me to update it? I’m now trying to fetch the updated features (DistTensor) in _sample_neighbors function before local_sample_neighbors.

github.com

dmlc/dgl/blob/master/python/dgl/distributed/graph_services.py#L63

    
      
              """ Sample from local partition.
          
          
    The input nodes use global IDs. We need to map the global node IDs to local node IDs,
              perform sampling and map the sampled results to the global IDs space again.
              The sampled results are stored in three vectors that store source nodes, destination nodes
              and edge IDs.
              """
              local_ids = partition_book.nid2localnid(seed_nodes, partition_book.partid)
              local_ids = F.astype(local_ids, local_g.idtype)
              # local_ids = self.seed_nodes
              sampled_graph = local_sample_neighbors(
                  local_g, local_ids, fan_out, edge_dir, prob, replace, _dist_training=True)
              global_nid_mapping = local_g.ndata[NID]
              src, dst = sampled_graph.edges()
              global_src, global_dst = F.gather_row(global_nid_mapping, src), \
                      F.gather_row(global_nid_mapping, dst)
              global_eids = F.gather_row(local_g.edata[EID], sampled_graph.edata[EID])
              return global_src, global_dst, global_eids
          
          
def _sample_etype_neighbors(local_g, partition_book, seed_nodes, etype_field,
                                      fan_out, edge_dir, prob, replace):

Is there any other suggestions?

VoVAllen · October 29, 2021, 10:44am

Yes prob is not available now but I think it’s possible to implement. One thing worth noticing is that the prob might need to be reweighted among different partitions. Image uniform sample edges with all one, if the partitions have different num edges, simply using the edge weight will make the prob differs for different partitions, which is not expected

VoVAllen · November 1, 2021, 7:14am

I just realized the prob is not affected by different partition. Please ignore my comment abouve

xcwanAndy · November 1, 2021, 7:43am

@VoVAllen Thanks a lot for your help. I’m implementing it now.

system · December 1, 2021, 7:44am

This topic was automatically closed 30 days after the last reply. New replies are no longer allowed.