dgl.sampling.select_topkο
- dgl.sampling.select_topk(g, k, weight, nodes=None, edge_dir='in', ascending=False, copy_ndata=True, copy_edata=True, output_device=None)[source]ο
Select the neighboring edges with k-largest (or k-smallest) weights of the given nodes and return the induced subgraph.
For each node, a number of inbound (or outbound when
edge_dir == 'out'
) edges with the largest (or smallest whenascending == True
) weights will be chosen. The graph returned will then contain all the nodes in the original graph, but only the sampled edges.Node/edge features are not preserved. The original IDs of the sampled edges are stored as the dgl.EID feature in the returned graph.
- Parameters:
g (DGLGraph) β The graph. Must be on CPU.
k (int or dict[etype, int]) β
The number of edges to be selected for each node on each edge type.
This argument can take a single int or a dictionary of edge types and ints. If a single int is given, DGL will select this number of edges for each node for every edge type.
If -1 is given for a single edge type, all the neighboring edges with that edge type will be selected.
weight (str) β Feature name of the weights associated with each edge. The feature should have only one element for each edge. The feature can be either int32/64 or float32/64.
nodes (tensor or dict, optional) β
Node IDs to sample neighbors from.
This argument can take a single ID tensor or a dictionary of node types and ID tensors. If a single tensor is given, the graph must only have one type of nodes.
If None, DGL will select the edges for all nodes.
edge_dir (str, optional) β
Determines whether to sample inbound or outbound edges.
Can take either
in
for inbound edges orout
for outbound edges.ascending (bool, optional) β If True, DGL will return edges with k-smallest weights instead of k-largest weights.
copy_ndata (bool, optional) β
If True, the node features of the new graph are copied from the original graph. If False, the new graph will not have any node features.
(Default: True)
copy_edata (bool, optional) β
If True, the edge features of the new graph are copied from the original graph. If False, the new graph will not have any edge features.
(Default: True)
output_device (Framework-specific device context object, optional) β The output device. Default is the same as the input graph.
- Returns:
A sampled subgraph containing only the sampled neighboring edges. It is on CPU.
- Return type:
Notes
If
copy_ndata
orcopy_edata
is True, same tensors are used as the node or edge features of the original graph and the new graph. As a result, users should avoid performing in-place operations on the node features of the new graph to avoid feature corruption.Examples
>>> g = dgl.graph(([0, 0, 1, 1, 2, 2], [1, 2, 0, 1, 2, 0])) >>> g.edata['weight'] = torch.FloatTensor([0, 1, 0, 1, 0, 1]) >>> sg = dgl.sampling.select_topk(g, 1, 'weight') >>> sg.edges(order='eid') (tensor([2, 1, 0]), tensor([0, 1, 2]))