Sampling function. Parameters ---------- g : DGLGraph The graph to sample nodes from. seed_nodes : Tensor or dict[str, Tensor] The nodes sampled in the current minibatch. exclude_eids : Tensor or dict[etype, Tensor], optional
(
self, g, seed_nodes, exclude_eids=None
)
| 88 | self.output_device = output_device |
| 89 | |
| 90 | def sample( |
| 91 | self, g, seed_nodes, exclude_eids=None |
| 92 | ): # pylint: disable=arguments-differ |
| 93 | """Sampling function. |
| 94 | |
| 95 | Parameters |
| 96 | ---------- |
| 97 | g : DGLGraph |
| 98 | The graph to sample nodes from. |
| 99 | seed_nodes : Tensor or dict[str, Tensor] |
| 100 | The nodes sampled in the current minibatch. |
| 101 | exclude_eids : Tensor or dict[etype, Tensor], optional |
| 102 | The edges to exclude from neighborhood expansion. |
| 103 | |
| 104 | Returns |
| 105 | ------- |
| 106 | input_nodes, output_nodes, subg |
| 107 | A triplet containing (1) the node IDs inducing the subgraph, (2) the node |
| 108 | IDs that are sampled in this minibatch, and (3) the subgraph itself. |
| 109 | """ |
| 110 | output_nodes = seed_nodes |
| 111 | for fanout in reversed(self.fanouts): |
| 112 | frontier = g.sample_neighbors( |
| 113 | seed_nodes, |
| 114 | fanout, |
| 115 | output_device=self.output_device, |
| 116 | replace=self.replace, |
| 117 | prob=self.prob, |
| 118 | exclude_edges=exclude_eids, |
| 119 | ) |
| 120 | block = transforms.to_block(frontier, seed_nodes) |
| 121 | seed_nodes = block.srcdata[NID] |
| 122 | |
| 123 | subg = g.subgraph( |
| 124 | seed_nodes, relabel_nodes=True, output_device=self.output_device |
| 125 | ) |
| 126 | if exclude_eids is not None: |
| 127 | subg = EidExcluder(exclude_eids)(subg) |
| 128 | |
| 129 | set_node_lazy_features(subg, self.prefetch_node_feats) |
| 130 | set_edge_lazy_features(subg, self.prefetch_edge_feats) |
| 131 | |
| 132 | return seed_nodes, output_nodes, subg |
nothing calls this directly
no test coverage detected