Method sample

python/dgl/dataloading/capped_neighbor_sampler.py:65–190

Sampling function. Parameters ---------- g : DGLGraph The graph to sample from. indices : Tensor or dict[str, Tensor] Nodes which induce the subgraph. exclude_eids : Tensor or dict[etype, Tensor], optional The edges to excl

(
        self, g, indices, exclude_eids=None
    )

        self.output_device = output_device

    def sample(
        self, g, indices, exclude_eids=None
    ):  # pylint: disable=arguments-differ
        """Sampling function.

        Parameters
        ----------
        g : DGLGraph
            The graph to sample from.
        indices : Tensor or dict[str, Tensor]
            Nodes which induce the subgraph.
        exclude_eids : Tensor or dict[etype, Tensor], optional
            The edges to exclude from the sampled subgraph.

        Returns
        -------
        input_nodes : Tensor or dict[str, Tensor]
            The node IDs inducing the subgraph.
        output_nodes : Tensor or dict[str, Tensor]
            The node IDs that are sampled in this minibatch.
        subg : DGLGraph
            The subgraph itself.
        """

        # Define empty dictionary to store reached nodes.
        output_nodes = indices
        all_reached_nodes = [indices]

        # Iterate over fanout.
        for fanout in reversed(self.fanouts):

            # Sample frontier.
            frontier = g.sample_neighbors(
                indices,
                fanout,
                output_device=self.output_device,
                replace=self.replace,
                prob=self.prob,
                exclude_edges=exclude_eids,
            )

            # Get reached nodes.
            curr_reached = defaultdict(list)
            for c_etype in frontier.canonical_etypes:
                (src_type, _, _) = c_etype
                src, _ = frontier.edges(etype=c_etype)
                curr_reached[src_type].append(src)

            # De-duplication.
            curr_reached = {
                ntype: torch.unique(torch.cat(srcs))
                for ntype, srcs in curr_reached.items()
            }

            # Generate type sampling probabilties.
            type_count = {
                node_type: indices.shape[0]
                for node_type, indices in curr_reached.items()