Repository: pyg-team/pytorch_geometric

Issue:
`self.training` is not passed to dropout layer after `to_hetero()` converts model
### 🐛 Describe the bug

Even when a user runs `model.eval()`, the model runs the dropout layer still with `training=True` if the model is converted with `to_hetero(model)`.

```python
import torch
import torch.nn.functional as F

from torch_geometric.nn import SAGEConv, to_hetero
from torch_geometric.data import HeteroData


class GNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = SAGEConv((-1, -1), 1)

    def forward(self, x, edge_index):
        x = self.conv(x, edge_index)
        return F.dropout(x, p=0.5, training=self.training)


def main():
    data = HeteroData()
    data["paper"].x = torch.randn(100, 3)
    data["author"].x = torch.randn(100, 5)
    data["author", "writes", "paper"].edge_index = torch.randint(0, 100, (2, 100))
    data["paper", "written", "author"].edge_index = torch.randint(0, 100, (2, 100))
    model = GNN()
    model = to_hetero(model, data.metadata(), aggr="sum")
    model.graph.print_tabular()

    model.eval()
    out1 = model(data.x_dict, data.edge_index_dict)
    out2 = model(data.x_dict, data.edge_index_dict)
    assert all(torch.allclose(o1, o2) for o1, o2 in zip(out1.values(), out2.values()))


if __name__ == "__main__":
    main()
```

```
opcode         name                                target                                 args                                                         kwargs
-------------  ----------------------------------  -------------------------------------  -----------------------------------------------------------  ----------------------------------------------
placeholder    x                                   x                                      ()                                                           {}
call_function  x_dict                              <function get_dict at 0x7f1a4ab8dea0>  (x,)                                                         {}
call_method    x__paper                            get                                    (x_dict, 'paper', None)                                      {}
call_method    x__author                           get                                    (x_dict, 'author', None)                                     {}
placeholder    edge_index                          edge_index                             ()                                                           {}
call_function  edge_index_dict                     <function get_dict at 0x7f1a4ab8dea0>  (edge_index,)                                                {}
call_method    edge_index__author__writes__paper   get                                    (edge_index_dict, ('author', 'writes', 'paper'), None)       {}
call_method    edge_index__paper__written__author  get                                    (edge_index_dict, ('paper', 'written', 'author'), None)      {}
call_module    conv__paper                         conv.author__writes__paper             ((x__author, x__paper), edge_index__author__writes__paper)   {}
call_module    conv__author                        conv.paper__written__author            ((x__paper, x__author), edge_index__paper__written__author)  {}
call_function  dropout__paper                      <function dropout at 0x7f1a61582170>   (conv__paper,)                                               {'p': 0.5, 'training': True, 'inplace': False}
call_function  dropout__author                     <function dropout at 0x7f1a61582170>   (conv__author,)                                              {'p': 0.5, 'training': True, 'inplace': False}
output         output                              output                                 ({'paper': dropout__paper, 'author': dropout__author},)      {}
Traceback (most recent call last):
  File "/workspaces/pytorch_geometric/test.py", line 34, in <module>
    main()
  File "/workspaces/pytorch_geometric/test.py", line 30, in main
    assert all(torch.allclose(o1, o2) for o1, o2 in zip(out1.values(), out2.values()))
AssertionError
```

### Environment

* PyG version: 9bc701731
* PyTorch version: 2.0.1
* OS: Linux codespaces-dc5000 5.15.0-1041-azure #48-Ubuntu SMP Tue Jun 20 20:34:08 UTC 2023 x86_64 x86_64 x86_64 GNU/Linux
* Python version: 3.10.8
* CUDA/cuDNN version: n/a
* How you installed PyTorch and PyG (`conda`, `pip`, source):
* Any other relevant information (*e.g.*, version of `torch-scatter`):



Discussion:
This issue arises only when a user has `F.dropout(..., training=self.training)` in their model directly because `to_hetero` doesn't touch submodules of the user's model at all:
https://github.com/pyg-team/pytorch_geometric/blob/165b68eebdc4d5bf44e56ca12f081aa55020a883/torch_geometric/nn/fx.py#L274-L276

This issue stems from `torch.fx.symbolic_trace` being unable to preserve the reference to `self.training`. During symbolic tracing, it sets a bool constant as the argument of `torch.nn.functional.dropout(..., training=True|False)`.

### Workaround
One way to avoid this issue is to replace the functional form `torch.nn.functional.dropout(...)` with its class form `torch.nn.Dropout(...)`

```diff
 class GNN(torch.nn.Module):
     def __init__(self):
         super().__init__()
         ...
+    self.dropout = torch.nn.Dropout(p=0.5)
     def forward(self, x, edge_index):
         ...
-        return F.dropout(x, p=0.5, training=self.training)
+        return self.dropout(x)
```

However, I cannot think of a way to fix this issue on PyG side so that all existing user code is automatically fixed in a newer PyG release.

References:
- https://pytorch.org/docs/2.0/fx.html#miscellanea
- https://pytorch.org/docs/2.0/fx.html#dynamic-control-flow
Thank you for this great catch. Shall we add a warning for this in `to_hetero`? Furthermore, is there any other workaround for this from PyTorch side (except using the module instance)?

Resolving patch:
diff --git a/torch_geometric/nn/fx.py b/torch_geometric/nn/fx.py
--- a/torch_geometric/nn/fx.py
+++ b/torch_geometric/nn/fx.py
@@ -1,4 +1,5 @@
 import copy
+import warnings
 from typing import Any, Dict, Optional
 
 import torch
@@ -127,6 +128,13 @@ def transform(self) -> GraphModule:
         # We iterate over each node and determine its output level
         # (node-level, edge-level) by filling `self._state`:
         for node in list(self.graph.nodes):
+            if node.op == 'call_function' and 'training' in node.kwargs:
+                warnings.warn(f"Found function '{node.name}' with keyword "
+                              f"argument 'training'. During FX tracing, this "
+                              f"will likely be baked in as a constant value. "
+                              f"Consider replacing this function by a module "
+                              f"to properly encapsulate its training flag.")
+
             if node.op == 'placeholder':
                 if node.name not in self._state:
                     if 'edge' in node.name or 'adj' in node.name:
diff --git a/torch_geometric/nn/models/basic_gnn.py b/torch_geometric/nn/models/basic_gnn.py
--- a/torch_geometric/nn/models/basic_gnn.py
+++ b/torch_geometric/nn/models/basic_gnn.py
@@ -2,7 +2,6 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
-import torch.nn.functional as F
 from torch import Tensor
 from torch.nn import Linear, ModuleList
 from tqdm import tqdm
@@ -83,7 +82,7 @@ def __init__(
         self.hidden_channels = hidden_channels
         self.num_layers = num_layers
 
-        self.dropout = dropout
+        self.dropout = torch.nn.Dropout(p=dropout)
         self.act = activation_resolver(act, **(act_kwargs or {}))
         self.jk_mode = jk
         self.act_first = act_first
@@ -232,7 +231,7 @@ def forward(
                 x = self.norms[i](x)
             if self.act is not None and not self.act_first:
                 x = self.act(x)
-            x = F.dropout(x, p=self.dropout, training=self.training)
+            x = self.dropout(x)
             if hasattr(self, 'jk'):
                 xs.append(x)
 
@@ -537,7 +536,7 @@ def init_conv(self, in_channels: Union[int, Tuple[int, int]],
 
         Conv = GATConv if not v2 else GATv2Conv
         return Conv(in_channels, out_channels, heads=heads, concat=concat,
-                    dropout=self.dropout, **kwargs)
+                    dropout=self.dropout.p, **kwargs)
 
 
 class PNA(BasicGNN):


Test patch:
diff --git a/test/nn/test_fx.py b/test/nn/test_fx.py
new file mode 100644
--- /dev/null
+++ b/test/nn/test_fx.py
@@ -0,0 +1,22 @@
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+def test_dropout():
+    class MyModule(torch.nn.Module):
+        def forward(self, x: Tensor) -> Tensor:
+            return F.dropout(x, p=1.0, training=self.training)
+
+    module = MyModule()
+    graph_module = torch.fx.symbolic_trace(module)
+    graph_module.recompile()
+
+    x = torch.randn(4)
+
+    graph_module.train()
+    assert torch.allclose(graph_module(x), torch.zeros_like(x))
+
+    # This is certainly undesired behavior due to tracing :(
+    graph_module.eval()
+    assert torch.allclose(graph_module(x), torch.zeros_like(x))
diff --git a/test/nn/test_model_summary.py b/test/nn/test_model_summary.py
--- a/test/nn/test_model_summary.py
+++ b/test/nn/test_model_summary.py
@@ -59,6 +59,7 @@ def test_summary_basic(gcn):
 | Layer               | Input Shape        | Output Shape   | #Param   |
 |---------------------+--------------------+----------------+----------|
 | GCN                 | [100, 32], [2, 20] | [100, 32]      | 1,072    |
+| ├─(dropout)Dropout  | [100, 16]          | [100, 16]      | --       |
 | ├─(act)ReLU         | [100, 16]          | [100, 16]      | --       |
 | ├─(convs)ModuleList | --                 | --             | 1,072    |
 | │    └─(0)GCNConv   | [100, 32], [2, 20] | [100, 16]      | 528      |
@@ -75,6 +76,7 @@ def test_summary_with_sparse_tensor(gcn):
 | Layer               | Input Shape           | Output Shape   | #Param   |
 |---------------------+-----------------------+----------------+----------|
 | GCN                 | [100, 32], [100, 100] | [100, 32]      | 1,072    |
+| ├─(dropout)Dropout  | [100, 16]             | [100, 16]      | --       |
 | ├─(act)ReLU         | [100, 16]             | [100, 16]      | --       |
 | ├─(convs)ModuleList | --                    | --             | 1,072    |
 | │    └─(0)GCNConv   | [100, 32], [100, 100] | [100, 16]      | 528      |
@@ -91,6 +93,7 @@ def test_summary_with_max_depth(gcn):
 | Layer               | Input Shape        | Output Shape   | #Param   |
 |---------------------+--------------------+----------------+----------|
 | GCN                 | [100, 32], [2, 20] | [100, 32]      | 1,072    |
+| ├─(dropout)Dropout  | [100, 16]          | [100, 16]      | --       |
 | ├─(act)ReLU         | [100, 16]          | [100, 16]      | --       |
 | ├─(convs)ModuleList | --                 | --             | 1,072    |
 +---------------------+--------------------+----------------+----------+
@@ -106,6 +109,7 @@ def test_summary_with_leaf_module(gcn):
 | Layer                                   | Input Shape        | Output Shape   | #Param   |
 |-----------------------------------------+--------------------+----------------+----------|
 | GCN                                     | [100, 32], [2, 20] | [100, 32]      | 1,072    |
+| ├─(dropout)Dropout                      | [100, 16]          | [100, 16]      | --       |
 | ├─(act)ReLU                             | [100, 16]          | [100, 16]      | --       |
 | ├─(convs)ModuleList                     | --                 | --             | 1,072    |
 | │    └─(0)GCNConv                       | [100, 32], [2, 20] | [100, 16]      | 528      |
diff --git a/test/nn/test_to_hetero_transformer.py b/test/nn/test_to_hetero_transformer.py
--- a/test/nn/test_to_hetero_transformer.py
+++ b/test/nn/test_to_hetero_transformer.py
@@ -277,7 +277,8 @@ def test_to_hetero_basic():
     assert out['author'].size() == (8, 16)
 
     model = Net10()
-    model = to_hetero(model, metadata, debug=False)
+    with pytest.warns(UserWarning, match="with keyword argument 'training'"):
+        model = to_hetero(model, metadata, debug=False)
     out = model(x_dict, edge_index_dict)
     assert isinstance(out, dict) and len(out) == 2
     assert out['paper'].size() == (100, 32)