Add dropout_node to torch_geometric.utils (#5481)

* add dropout_node

* pass num_nodes to subgraph

* drop relabel_nodes argument

* Update torch_geometric/utils/dropout.py

Co-authored-by: Matthias Fey <[email protected]>

* Update torch_geometric/utils/dropout.py

Co-authored-by: Matthias Fey <[email protected]>

* Update torch_geometric/utils/dropout.py

Co-authored-by: Matthias Fey <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update

* test

* changelog

Co-authored-by: Matthias Fey <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [2.2.0] - 2022-MM-DD
 ### Added
+- Added `dropout_node` augmentation that randomly drops nodes from a graph ([#5481](https://github.com/pyg-team/pytorch_geometric/pull/5481))
 - Added `AddRandomMetaPaths` that adds edges based on random walks along a metapath ([#5397](https://github.com/pyg-team/pytorch_geometric/pull/5397))
 - Added `WLConvContinuous` for performing WL refinement with continuous attributes ([#5316](https://github.com/pyg-team/pytorch_geometric/pull/5316))
 - Added `print_summary` method for the `torch_geometric.data.Dataset` interface ([#5438](https://github.com/pyg-team/pytorch_geometric/pull/5438))

test/utils/test_dropout.py

@@ -1,6 +1,6 @@
 import torch
 
-from torch_geometric.utils import dropout_adj
+from torch_geometric.utils import dropout_adj, dropout_node
 
 
 def test_dropout_adj():
@@ -23,3 +23,21 @@ def test_dropout_adj():
     out = dropout_adj(edge_index, edge_attr, force_undirected=True)
     assert out[0].tolist() == [[0, 1, 1, 2], [1, 2, 0, 1]]
     assert out[1].tolist() == [1, 3, 1, 3]
+
+
+def test_dropout_node():
+    edge_index = torch.tensor([
+        [0, 1, 1, 2, 2, 3],
+        [1, 0, 2, 1, 3, 2],
+    ])
+
+    out = dropout_node(edge_index, training=False)
+    assert edge_index.tolist() == out[0].tolist()
+    assert out[1].tolist() == [False, False, False, False, False, False]
+    assert out[2].tolist() == [False, False, False, False]
+
+    torch.manual_seed(5)
+    out = dropout_node(edge_index)
+    assert out[0].tolist() == [[2, 3], [3, 2]]
+    assert out[1].tolist() == [False, False, False, False, True, True]
+    assert out[2].tolist() == [True, False, True, True]

torch_geometric/utils/__init__.py

@@ -1,6 +1,6 @@
 from .degree import degree
 from .softmax import softmax
-from .dropout import dropout_adj
+from .dropout import dropout_adj, dropout_node
 from .sort_edge_index import sort_edge_index
 from .coalesce import coalesce
 from .undirected import is_undirected, to_undirected
@@ -39,6 +39,7 @@
     'degree',
     'softmax',
     'dropout_adj',
+    'dropout_node',
     'sort_edge_index',
     'coalesce',
     'is_undirected',

torch_geometric/utils/dropout.py

@@ -5,6 +5,9 @@
 
 from torch_geometric.typing import OptTensor
 
+from .num_nodes import maybe_num_nodes
+from .subgraph import subgraph
+
 
 def filter_adj(row: Tensor, col: Tensor, edge_attr: OptTensor,
                mask: Tensor) -> Tuple[Tensor, Tensor, OptTensor]:
@@ -79,3 +82,56 @@ def dropout_adj(
         edge_index = torch.stack([row, col], dim=0)
 
     return edge_index, edge_attr
+
+
+def dropout_node(edge_index: Tensor, p: float = 0.5,
+                 num_nodes: Optional[int] = None,
+                 training: bool = True) -> Tuple[Tensor, Tensor, Tensor]:
+    r"""Randomly drops nodes from the adjacency matrix
+    :obj:`edge_index` with probability :obj:`p` using samples from
+    a Bernoulli distribution.
+
+    The method returns (1) the retained :obj:`edge_index`, (2) the edge mask
+    indicating which edges were dropped. (3) the node mask indicating
+    which nodes were dropped.
+
+    Args:
+        edge_index (LongTensor): The edge indices.
+        p (float, optional): Dropout probability. (default: :obj:`0.5`)
+        num_nodes (int, optional): The number of nodes, *i.e.*
+            :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
+        training (bool, optional): If set to :obj:`False`, this operation is a
+            no-op. (default: :obj:`True`)
+
+    :rtype: (:class:`LongTensor`, :class:`BoolTensor`, :class:`BoolTensor`)
+
+    Examples:
+
+        >>> edge_index = torch.tensor([[0, 1, 1, 2, 2, 3],
+        ...                            [1, 0, 2, 1, 3, 2]])
+        >>> edge_index, edge_mask, node_mask = dropout_node(edge_index)
+        >>> edge_index
+        tensor([[0, 1],
+                [1, 0]])
+        >>> edge_mask
+        tensor([ True,  True, False, False, False, False])
+        >>> node_mask
+        tensor([ True,  True, False, False])
+    """
+    if p < 0. or p > 1.:
+        raise ValueError(f'Dropout probability has to be between 0 and 1 '
+                         f'(got {p}')
+
+    num_nodes = maybe_num_nodes(edge_index, num_nodes)
+
+    if not training or p == 0.0:
+        node_mask = edge_index.new_zeros(num_nodes, dtype=torch.bool)
+        edge_mask = edge_index.new_zeros(edge_index.size(1), dtype=torch.bool)
+        return edge_index, edge_mask, node_mask
+
+    prob = torch.rand(num_nodes, device=edge_index.device)
+    node_mask = prob > p
+    edge_index, _, edge_mask = subgraph(node_mask, edge_index,
+                                        num_nodes=num_nodes,
+                                        return_edge_mask=True)
+    return edge_index, edge_mask, node_mask