Repository: pyg-team/pytorch_geometric (python)

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`):


Pull request description:
Warn user when using `training` flag in `to_hetero`
Fixes #7745

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9bcc1c405..d0436ca8e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -74,6 +74,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Warn user when using the `training` flag in `to_hetero` modules ([#7772](https://github.com/pyg-team/pytorch_geometric/pull/7772))
 - Unchained exceptions raised when accessing non-existent data attributes for better readability ([#7734](https://github.com/pyg-team/pytorch_geometric/pull/7734))
 - Raise error when collecting non-existing attributes in `HeteroData` ([#7714](https://github.com/pyg-team/pytorch_geometric/pull/7714))
 - Renamed `dest` argument to `dst` in `utils.geodesic_distance` ([#7708](https://github.com/pyg-team/pytorch_geometric/pull/7708))
diff --git a/torch_geometric/data/hetero_data.py b/torch_geometric/data/hetero_data.py
index cfe58a3b2..0e2955d7b 100644
--- a/torch_geometric/data/hetero_data.py
+++ b/torch_geometric/data/hetero_data.py
@@ -331,7 +331,7 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
                     **kwargs) -> Any:
         if isinstance(value, SparseTensor) and 'adj' in key:
             return (0, 1)
-        elif 'index' in key or 'face' in key:
+        elif isinstance(store, EdgeStorage) and 'index' in key:
             return -1
         return 0
 
diff --git a/torch_geometric/nn/fx.py b/torch_geometric/nn/fx.py
index 860def82a..021f43ca6 100644
--- 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 @@ class Transformer:
         # 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
index 7244893fc..15397a6b5 100644
--- a/torch_geometric/nn/models/basic_gnn.py
+++ b/torch_geometric/nn/models/basic_gnn.py
@@ -2,7 +2,6 @@ import copy
 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 @@ class BasicGNN(torch.nn.Module):
         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 @@ class BasicGNN(torch.nn.Module):
                 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 @@ class GAT(BasicGNN):
 
         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
index 000000000..a0f1ff44c
--- /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
index 6408e0ca0..f07aff03c 100644
--- 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
index 768806bbd..6ee4ed316 100644
--- 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)