Method pin_memory_

python/dgl/heterograph.py:5773–5843

Pin the graph structure and node/edge data to the page-locked memory for GPU zero-copy access. This is an **inplace** method. The graph structure must be on CPU to be pinned. If the graph struture is already pinned, the function directly returns it. Materialization

(self)

        return self

    def pin_memory_(self):
        """Pin the graph structure and node/edge data to the page-locked memory for
        GPU zero-copy access.

        This is an **inplace** method. The graph structure must be on CPU to be pinned.
        If the graph struture is already pinned, the function directly returns it.

        Materialization of new sparse formats for pinned graphs is not allowed.
        To avoid implicit formats materialization during training,
        you should create all the needed formats before pinning.
        But cloning and materialization is fine. See the examples below.

        Returns
        -------
        DGLGraph
            The pinned graph.

        Examples
        --------
        The following example uses PyTorch backend.

        >>> import dgl
        >>> import torch

        >>> g = dgl.graph((torch.tensor([1, 0]), torch.tensor([1, 2])))
        >>> g.pin_memory_()

        Materialization of new sparse formats is not allowed for pinned graphs.

        >>> g.create_formats_()  # This would raise an error! You should do this before pinning.

        Cloning and materializing new formats is allowed. The returned graph is **not** pinned.

        >>> g1 = g.formats(['csc'])
        >>> assert not g1.is_pinned()

        The pinned graph can be access from both CPU and GPU. The concrete device depends
        on the context of ``query``. For example, ``eid`` in ``find_edges()`` is a query.
        When ``eid`` is on CPU, ``find_edges()`` is executed on CPU, and the returned
        values are CPU tensors

        >>> g.unpin_memory_()
        >>> g.create_formats_()
        >>> g.pin_memory_()
        >>> eid = torch.tensor([1])
        >>> g.find_edges(eids)
        (tensor([0]), tensor([2]))

        Moving ``eid`` to GPU, ``find_edges()`` will be executed on GPU, and the returned
        values are GPU tensors.

        >>> eid = eid.to('cuda:0')
        >>> g.find_edges(eids)
        (tensor([0], device='cuda:0'), tensor([2], device='cuda:0'))

        If you don't provide a ``query``, methods will be executed on CPU by default.

        >>> g.in_degrees()