Repository: pyg-team/pytorch_geometric
Pull request #7228 (2023-04-24 17:45:11)

minor fix for `qm9_pretrained_schnet.py`
```
root@9135585df2bb:/workspace# python3 examples/qm9_pretrained_schnet.py
Traceback (most recent call last):
  File "examples/qm9_pretrained_schnet.py", line 17, in <module>
    dataset = QM9(osp.join())
TypeError: join() missing 1 required positional argument: 'a'
```

after this minor fix and pip installing ase and schnet==1.0.0
i get:
```
Traceback (most recent call last):
  File "examples/qm9_pretrained_schnet.py", line 22, in <module>
    model, datasets = SchNet.from_qm9_pretrained(path, dataset, target)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric/nn/models/schnet.py", line 252, in from_qm9_pretrained
    net.readout = 'mean' if mean is True else 'add'
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1685, in __setattr__
    raise TypeError("cannot assign '{}' as child module '{}' "
TypeError: cannot assign 'str' as child module 'readout' (torch.nn.Module or None expected)
```
if i set net.readout = None then it works and I get to the part of the schnet that relies on torch cluster (which i will eventually make a part of pyg-lib along with a bunch of other torch-* functionalities that are needed to be moved to pyg-lib)

```
  File "examples/qm9_pretrained_schnet.py", line 32, in <module>
    pred = model(data.z, data.pos, data.batch)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1533, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric/nn/models/schnet.py", line 284, in forward
    edge_index, edge_weight = self.interaction_graph(pos, batch)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1533, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric/nn/models/schnet.py", line 352, in forward
    edge_index = radius_graph(pos, r=self.cutoff, batch=batch,
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric/nn/pool/__init__.py", line 210, in radius_graph
    return torch_cluster.radius_graph(x, r, batch, loop, max_num_neighbors,
AttributeError: 'NoneType' object has no attribute 'radius_graph'
```

note that w/ schnet >= 2.0:
```
Traceback (most recent call last):
  File "examples/qm9_pretrained_schnet.py", line 22, in <module>
    model, datasets = SchNet.from_qm9_pretrained(path, dataset, target)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric/nn/models/schnet.py", line 219, in from_qm9_pretrained
    state = torch.load(path, map_location='cpu')
  File "/usr/local/lib/python3.8/dist-packages/torch/serialization.py", line 817, in load
    return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args)
  File "/usr/local/lib/python3.8/dist-packages/torch/serialization.py", line 1045, in _legacy_load
    result = unpickler.load()
  File "/usr/lib/python3.8/pickle.py", line 1212, in load
    dispatch[key[0]](self)
  File "/usr/lib/python3.8/pickle.py", line 1528, in load_global
    klass = self.find_class(module, name)
  File "/usr/local/lib/python3.8/dist-packages/torch/serialization.py", line 850, in find_class
    return super().find_class(mod_name, name)
  File "/usr/local/lib/python3.8/dist-packages/pytorch_lightning/_graveyard/legacy_import_unpickler.py", line 24, in find_class
    return super().find_class(new_module, name)
  File "/usr/lib/python3.8/pickle.py", line 1579, in find_class
    __import__(module, level=0)
ModuleNotFoundError: No module named 'schnetpack.atomistic.model'
```




Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 03eaaec8b..be6c4b675 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -29,6 +29,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Fixed `HGTConv` utility function `_construct_src_node_feat` ([#7194](https://github.com/pyg-team/pytorch_geometric/pull/7194))
 - Extend dataset summary to create stats for each node/edge type ([#7203](https://github.com/pyg-team/pytorch_geometric/pull/7203))
 - Added an optional `batch_size` argument to `avg_pool_x` and `max_pool_x` ([#7216](https://github.com/pyg-team/pytorch_geometric/pull/7216))
 - Fixed `subgraph` on unordered inputs ([#7187](https://github.com/pyg-team/pytorch_geometric/pull/7187))
diff --git a/examples/qm9_pretrained_schnet.py b/examples/qm9_pretrained_schnet.py
index da91f1db0..61d6c31bf 100644
--- a/examples/qm9_pretrained_schnet.py
+++ b/examples/qm9_pretrained_schnet.py
@@ -14,7 +14,7 @@ parser.add_argument('--cutoff', type=float, default=10.0,
 args = parser.parse_args()
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'QM9')
-dataset = QM9(osp.join())
+dataset = QM9(path)
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
diff --git a/torch_geometric/nn/conv/hgt_conv.py b/torch_geometric/nn/conv/hgt_conv.py
index a97f2121a..9287f6501 100644
--- a/torch_geometric/nn/conv/hgt_conv.py
+++ b/torch_geometric/nn/conv/hgt_conv.py
@@ -71,6 +71,10 @@ class HGTConv(MessagePassing):
         self.heads = heads
         self.node_types = metadata[0]
         self.edge_types = metadata[1]
+        self.edge_types_map = {
+            edge_type: i
+            for i, edge_type in enumerate(metadata[1])
+        }
 
         self.dst_node_types = set([key[-1] for key in self.edge_types])
 
@@ -83,10 +87,10 @@ class HGTConv(MessagePassing):
         dim = out_channels // heads
         num_types = heads * len(self.edge_types)
 
-        self.k_rel = HeteroLinear(dim, dim, num_types, is_sorted=True,
-                                  bias=False)
-        self.v_rel = HeteroLinear(dim, dim, num_types, is_sorted=True,
-                                  bias=False)
+        self.k_rel = HeteroLinear(dim, dim, num_types, bias=False,
+                                  is_sorted=True)
+        self.v_rel = HeteroLinear(dim, dim, num_types, bias=False,
+                                  is_sorted=True)
 
         self.skip = ParameterDict({
             node_type: Parameter(torch.Tensor(1))
@@ -121,36 +125,40 @@ class HGTConv(MessagePassing):
         return torch.cat(outs, dim=0), offset
 
     def _construct_src_node_feat(
-        self,
-        k_dict: Dict[str, Tensor],
-        v_dict: Dict[str, Tensor],
+        self, k_dict: Dict[str, Tensor], v_dict: Dict[str, Tensor],
+        edge_index_dict: Dict[EdgeType, Adj]
     ) -> Tuple[Tensor, Tensor, Dict[EdgeType, int]]:
         """Constructs the source node representations."""
-        count = 0
         cumsum = 0
+        num_edge_types = len(self.edge_types)
         H, D = self.heads, self.out_channels // self.heads
 
         # Flatten into a single tensor with shape [num_edge_types * heads, D]:
         ks: List[Tensor] = []
         vs: List[Tensor] = []
-        type_list: List[int] = []
+        type_list: List[Tensor] = []
         offset: Dict[EdgeType] = {}
-        for edge_type in self.edge_types:
-            src, _, _ = edge_type
-
-            ks.append(k_dict[src].reshape(-1, D))
-            vs.append(v_dict[src].reshape(-1, D))
-
+        for edge_type in edge_index_dict.keys():
+            src = edge_type[0]
             N = k_dict[src].size(0)
-            for _ in range(H):
-                type_list.append(torch.full((N, ), count, dtype=torch.long))
-                count += 1
             offset[edge_type] = cumsum
             cumsum += N
 
-        type_vec = torch.cat(type_list, dim=0)
-        k = self.k_rel(torch.cat(ks, dim=0), type_vec).view(-1, H, D)
-        v = self.v_rel(torch.cat(vs, dim=0), type_vec).view(-1, H, D)
+            # construct type_vec for curr edge_type with shape [H, D]
+            edge_type_offset = self.edge_types_map[edge_type]
+            type_vec = torch.arange(H, dtype=torch.long).view(-1, 1).repeat(
+                1, N) * num_edge_types + edge_type_offset
+
+            type_list.append(type_vec)
+            ks.append(k_dict[src])
+            vs.append(v_dict[src])
+
+        ks = torch.cat(ks, dim=0).transpose(0, 1).reshape(-1, D)
+        vs = torch.cat(vs, dim=0).transpose(0, 1).reshape(-1, D)
+        type_vec = torch.cat(type_list, dim=1).flatten()
+
+        k = self.k_rel(ks, type_vec).view(H, -1, D).transpose(0, 1)
+        v = self.v_rel(vs, type_vec).view(H, -1, D).transpose(0, 1)
 
         return k, v, offset
 
@@ -184,12 +192,14 @@ class HGTConv(MessagePassing):
         # Compute K, Q, V over node types:
         kqv_dict = self.kqv_lin(x_dict)
         for key, val in kqv_dict.items():
-            k_dict[key] = val[:, :F].view(-1, H, D)
-            q_dict[key] = val[:, F:2 * F].view(-1, H, D)
-            v_dict[key] = val[:, 2 * F:].view(-1, H, D)
+            k, q, v = torch.tensor_split(val, 3, dim=1)
+            k_dict[key] = k.view(-1, H, D)
+            q_dict[key] = q.view(-1, H, D)
+            v_dict[key] = v.view(-1, H, D)
 
         q, dst_offset = self._cat(q_dict)
-        k, v, src_offset = self._construct_src_node_feat(k_dict, v_dict)
+        k, v, src_offset = self._construct_src_node_feat(
+            k_dict, v_dict, edge_index_dict)
 
         edge_index, edge_attr = construct_bipartite_edge_index(
             edge_index_dict, src_offset, dst_offset, edge_attr_dict=self.p_rel)
@@ -200,7 +210,8 @@ class HGTConv(MessagePassing):
         # Reconstruct output node embeddings dict:
         for node_type, start_offset in dst_offset.items():
             end_offset = start_offset + q_dict[node_type].size(0)
-            out_dict[node_type] = out[start_offset:end_offset]
+            if node_type in self.dst_node_types:
+                out_dict[node_type] = out[start_offset:end_offset]
 
         # Transform output node embeddings:
         a_dict = self.out_lin({
@@ -210,11 +221,7 @@ class HGTConv(MessagePassing):
 
         # Iterate over node types:
         for node_type, out in out_dict.items():
-            if node_type not in self.dst_node_types:
-                out_dict[node_type] = None
-                continue
-            else:
-                out = a_dict[node_type]
+            out = a_dict[node_type]
 
             if out.size(-1) == x_dict[node_type].size(-1):
                 alpha = self.skip[node_type].sigmoid()
diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py
index e042abb9e..b13c474bd 100644
--- a/torch_geometric/nn/dense/linear.py
+++ b/torch_geometric/nn/dense/linear.py
@@ -387,7 +387,7 @@ class HeteroDictLinear(torch.nn.Module):
                     biases.append(lin.bias)
             biases = None if biases[0] is None else biases
             outs = pyg_lib.ops.grouped_matmul(xs, weights, biases)
-            for key, out in zip(self.lins.keys(), outs):
+            for key, out in zip(x_dict.keys(), outs):
                 if key in x_dict:
                     out_dict[key] = out
         else:


Test patch:
diff --git a/test/nn/conv/test_hgt_conv.py b/test/nn/conv/test_hgt_conv.py
index 02bb7c8e4..a41a755a9 100644
--- a/test/nn/conv/test_hgt_conv.py
+++ b/test/nn/conv/test_hgt_conv.py
@@ -193,7 +193,45 @@ def test_hgt_conv_missing_dst_node_type():
     out_dict = conv(data.x_dict, data.edge_index_dict)
     assert out_dict['author'].size() == (4, 64)
     assert out_dict['paper'].size() == (6, 64)
-    assert out_dict['university'] is None
+    assert 'university' not in out_dict
+
+
+def test_hgt_conv_missing_input_node_type():
+    data = HeteroData()
+    data['author'].x = torch.randn(4, 16)
+    data['paper'].x = torch.randn(6, 32)
+    data['author', 'writes',
+         'paper'].edge_index = get_random_edge_index(4, 6, 20)
+
+    # Some nodes from metadata are missing in data.
+    # This might happen while using NeighborLoader.
+    metadata = (['author', 'paper',
+                 'university'], [('author', 'writes', 'paper')])
+    conv = HGTConv(-1, 64, metadata, heads=1)
+
+    out_dict = conv(data.x_dict, data.edge_index_dict)
+    assert out_dict['paper'].size() == (6, 64)
+    assert 'university' not in out_dict
+
+
+def test_hgt_conv_missing_edge_type():
+    data = HeteroData()
+    data['author'].x = torch.randn(4, 16)
+    data['paper'].x = torch.randn(6, 32)
+    data['university'].x = torch.randn(10, 32)
+
+    data['author', 'writes',
+         'paper'].edge_index = get_random_edge_index(4, 6, 20)
+
+    metadata = (['author', 'paper',
+                 'university'], [('author', 'writes', 'paper'),
+                                 ('university', 'employs', 'author')])
+    conv = HGTConv(-1, 64, metadata, heads=1)
+
+    out_dict = conv(data.x_dict, data.edge_index_dict)
+    assert out_dict['author'].size() == (4, 64)
+    assert out_dict['paper'].size() == (6, 64)
+    assert 'university' not in out_dict
 
 
 if __name__ == '__main__':