pyg-nightly 2.6.0.dev20240319__py3-none-any.whl → 2.7.0.dev20250114__py3-none-any.whl

torch_geometric/nn/nlp/sentence_transformer.py

@@ -0,0 +1,134 @@
+from enum import Enum
+from typing import List, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+
+class PoolingStrategy(Enum):
+    MEAN = 'mean'
+    LAST = 'last'
+    CLS = 'cls'
+    LAST_HIDDEN_STATE = 'last_hidden_state'
+
+
+class SentenceTransformer(torch.nn.Module):
+    def __init__(
+        self,
+        model_name: str,
+        pooling_strategy: Union[PoolingStrategy, str] = 'mean',
+    ) -> None:
+        super().__init__()
+
+        self.model_name = model_name
+        self.pooling_strategy = PoolingStrategy(pooling_strategy)
+
+        from transformers import AutoModel, AutoTokenizer
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+    def forward(self, input_ids: Tensor, attention_mask: Tensor) -> Tensor:
+        out = self.model(input_ids=input_ids, attention_mask=attention_mask)
+
+        emb = out[0]  # First element contains all token embeddings.
+        if self.pooling_strategy == PoolingStrategy.MEAN:
+            emb = mean_pooling(emb, attention_mask)
+        elif self.pooling_strategy == PoolingStrategy.LAST:
+            emb = last_pooling(emb, attention_mask)
+        elif self.pooling_strategy == PoolingStrategy.LAST_HIDDEN_STATE:
+            emb = out.last_hidden_state
+        else:
+            assert self.pooling_strategy == PoolingStrategy.CLS
+            emb = emb[:, 0, :]
+
+        emb = F.normalize(emb, p=2, dim=1)
+        return emb
+
+    def get_input_ids(
+        self,
+        text: List[str],
+        batch_size: Optional[int] = None,
+        output_device: Optional[Union[torch.device, str]] = None,
+    ) -> Tensor:
+        is_empty = len(text) == 0
+        text = ['dummy'] if is_empty else text
+
+        batch_size = len(text) if batch_size is None else batch_size
+
+        input_ids: List[Tensor] = []
+        attention_masks: List[Tensor] = []
+        for start in range(0, len(text), batch_size):
+            token = self.tokenizer(
+                text[start:start + batch_size],
+                padding=True,
+                truncation=True,
+                return_tensors='pt',
+            )
+            input_ids.append(token.input_ids.to(self.device))
+            attention_masks.append(token.attention_mask.to(self.device))
+
+        def _out(x: List[Tensor]) -> Tensor:
+            out = torch.cat(x, dim=0) if len(x) > 1 else x[0]
+            out = out[:0] if is_empty else out
+            return out.to(output_device)
+
+        return _out(input_ids), _out(attention_masks)
+
+    @property
+    def device(self) -> torch.device:
+        return next(iter(self.model.parameters())).device
+
+    @torch.no_grad()
+    def encode(
+        self,
+        text: List[str],
+        batch_size: Optional[int] = None,
+        output_device: Optional[Union[torch.device, str]] = None,
+    ) -> Tensor:
+        is_empty = len(text) == 0
+        text = ['dummy'] if is_empty else text
+
+        batch_size = len(text) if batch_size is None else batch_size
+
+        embs: List[Tensor] = []
+        for start in range(0, len(text), batch_size):
+            token = self.tokenizer(
+                text[start:start + batch_size],
+                padding=True,
+                truncation=True,
+                return_tensors='pt',
+            )
+
+            emb = self(
+                input_ids=token.input_ids.to(self.device),
+                attention_mask=token.attention_mask.to(self.device),
+            ).to(output_device)
+
+            embs.append(emb)
+
+        out = torch.cat(embs, dim=0) if len(embs) > 1 else embs[0]
+        out = out[:0] if is_empty else out
+        return out
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(model_name={self.model_name})'
+
+
+def mean_pooling(emb: Tensor, attention_mask: Tensor) -> Tensor:
+    mask = attention_mask.unsqueeze(-1).expand(emb.size()).to(emb.dtype)
+    return (emb * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
+
+
+def last_pooling(emb: Tensor, attention_mask: Tensor) -> Tensor:
+    # Check whether language model uses left padding,
+    # which is always used for decoder LLMs
+    left_padding = attention_mask[:, -1].sum() == attention_mask.size(0)
+    if left_padding:
+        return emb[:, -1]
+
+    seq_indices = attention_mask.sum(dim=1) - 1
+    return emb[torch.arange(emb.size(0), device=emb.device), seq_indices]

torch_geometric/nn/nlp/vision_transformer.py

@@ -0,0 +1,33 @@
+from typing import Optional, Union
+
+import torch
+from torch import Tensor
+
+
+class VisionTransformer(torch.nn.Module):
+    def __init__(
+        self,
+        model_name: str,
+    ) -> None:
+        super().__init__()
+        self.model_name = model_name
+
+        from transformers import SwinConfig, SwinModel
+
+        self.config = SwinConfig.from_pretrained(model_name)
+        self.model = SwinModel(self.config)
+
+    @torch.no_grad()
+    def forward(
+        self,
+        images: Tensor,
+        output_device: Optional[Union[torch.device, str]] = None,
+    ) -> Tensor:
+        return self.model(images).last_hidden_state.to(output_device)
+
+    @property
+    def device(self) -> torch.device:
+        return next(iter(self.model.parameters())).device
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(model_name={self.model_name})'

torch_geometric/nn/norm/batch_norm.py

@@ -88,7 +88,7 @@ class BatchNorm(torch.nn.Module):
         return self.module(x)
 
     def __repr__(self):
-        return f'{self.__class__.__name__}({self.module.num_features})'
+        return f'{self.__class__.__name__}({self.module.extra_repr()})'
 
 
 class HeteroBatchNorm(torch.nn.Module):

torch_geometric/nn/parameter_dict.py

@@ -1,4 +1,4 @@
-from typing import Final, Iterable, Mapping, Optional, Set, Tuple, Union
+from typing import Final, Iterable, Mapping, Optional, Tuple, Union
 
 import torch
 from torch.nn import Parameter
@@ -11,7 +11,7 @@ Key = Union[str, Tuple[str, ...]]
 # internal representation and converts it back to `.` in the external
 # representation. It also allows passing tuples as keys.
 class ParameterDict(torch.nn.ParameterDict):
-    CLASS_ATTRS: Final[Set[str]] = set(dir(torch.nn.ParameterDict))
+    CLASS_ATTRS: Final[Tuple[str, ...]] = set(dir(torch.nn.ParameterDict))
 
     def __init__(
         self,

torch_geometric/nn/pool/__init__.py

@@ -7,18 +7,19 @@ from torch import Tensor
 import torch_geometric.typing
 from torch_geometric.typing import OptTensor, torch_cluster
 
-from .asap import ASAPooling
 from .avg_pool import avg_pool, avg_pool_neighbor_x, avg_pool_x
-from .edge_pool import EdgePooling
 from .glob import global_add_pool, global_max_pool, global_mean_pool
 from .knn import (KNNIndex, L2KNNIndex, MIPSKNNIndex, ApproxL2KNNIndex,
                   ApproxMIPSKNNIndex)
 from .graclus import graclus
 from .max_pool import max_pool, max_pool_neighbor_x, max_pool_x
-from .mem_pool import MemPooling
-from .pan_pool import PANPooling
-from .sag_pool import SAGPooling
 from .topk_pool import TopKPooling
+from .sag_pool import SAGPooling
+from .edge_pool import EdgePooling
+from .cluster_pool import ClusterPooling
+from .asap import ASAPooling
+from .pan_pool import PANPooling
+from .mem_pool import MemPooling
 from .voxel_grid import voxel_grid
 from .approx_knn import approx_knn, approx_knn_graph
 
@@ -218,6 +219,13 @@ def radius(
             Automatically calculated if not given. (default: :obj:`None`)
 
     :rtype: :class:`torch.Tensor`
+
+    .. warning::
+
+        The CPU implementation of :meth:`radius` with :obj:`max_num_neighbors`
+        is biased towards certain quadrants.
+        Consider setting :obj:`max_num_neighbors` to :obj:`None` or moving
+        inputs to GPU before proceeding.
     """
     if not torch_geometric.typing.WITH_TORCH_CLUSTER_BATCH_SIZE:
         return torch_cluster.radius(x, y, r, batch_x, batch_y,
@@ -268,6 +276,13 @@ def radius_graph(
             Automatically calculated if not given. (default: :obj:`None`)
 
     :rtype: :class:`torch.Tensor`
+
+    .. warning::
+
+        The CPU implementation of :meth:`radius_graph` with
+        :obj:`max_num_neighbors` is biased towards certain quadrants.
+        Consider setting :obj:`max_num_neighbors` to :obj:`None` or moving
+        inputs to GPU before proceeding.
     """
     if batch is not None and x.device != batch.device:
         warnings.warn("Input tensor 'x' and 'batch' are on different devices "
@@ -330,6 +345,7 @@ __all__ = [
     'TopKPooling',
     'SAGPooling',
     'EdgePooling',
+    'ClusterPooling',
     'ASAPooling',
     'PANPooling',
     'MemPooling',

torch_geometric/nn/pool/cluster_pool.py

@@ -0,0 +1,145 @@
+from typing import NamedTuple, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from torch_geometric.utils import (
+    dense_to_sparse,
+    one_hot,
+    to_dense_adj,
+    to_scipy_sparse_matrix,
+)
+
+
+class UnpoolInfo(NamedTuple):
+    edge_index: Tensor
+    cluster: Tensor
+    batch: Tensor
+
+
+class ClusterPooling(torch.nn.Module):
+    r"""The cluster pooling operator from the `"Edge-Based Graph Component
+    Pooling" <paper url>`_ paper.
+
+    :class:`ClusterPooling` computes a score for each edge.
+    Based on the selected edges, graph clusters are calculated and compressed
+    to one node using the injective :obj:`"sum"` aggregation function.
+    Edges are remapped based on the nodes created by each cluster and the
+    original edges.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        edge_score_method (str, optional): The function to apply
+            to compute the edge score from raw edge scores (:obj:`"tanh"`,
+            :obj:`"sigmoid"`, :obj:`"log_softmax"`). (default: :obj:`"tanh"`)
+        dropout (float, optional): The probability with
+            which to drop edge scores during training. (default: :obj:`0.0`)
+        threshold (float, optional): The threshold of edge scores. If set to
+            :obj:`None`, will be automatically inferred depending on
+            :obj:`edge_score_method`. (default: :obj:`None`)
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        edge_score_method: str = 'tanh',
+        dropout: float = 0.0,
+        threshold: Optional[float] = None,
+    ):
+        super().__init__()
+        assert edge_score_method in ['tanh', 'sigmoid', 'log_softmax']
+
+        if threshold is None:
+            threshold = 0.5 if edge_score_method == 'sigmoid' else 0.0
+
+        self.in_channels = in_channels
+        self.edge_score_method = edge_score_method
+        self.dropout = dropout
+        self.threshhold = threshold
+
+        self.lin = torch.nn.Linear(2 * in_channels, 1)
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.lin.reset_parameters()
+
+    def forward(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+    ) -> Tuple[Tensor, Tensor, Tensor, UnpoolInfo]:
+        r"""Forward pass.
+
+        Args:
+            x (torch.Tensor): The node features.
+            edge_index (torch.Tensor): The edge indices.
+            batch (torch.Tensor): Batch vector
+                :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns
+                each node to a specific example.
+
+        Return types:
+            * **x** *(torch.Tensor)* - The pooled node features.
+            * **edge_index** *(torch.Tensor)* - The coarsened edge indices.
+            * **batch** *(torch.Tensor)* - The coarsened batch vector.
+            * **unpool_info** *(UnpoolInfo)* - Information that can be consumed
+              for unpooling.
+        """
+        mask = edge_index[0] != edge_index[1]
+        edge_index = edge_index[:, mask]
+
+        edge_attr = torch.cat(
+            [x[edge_index[0]], x[edge_index[1]]],
+            dim=-1,
+        )
+        edge_score = self.lin(edge_attr).view(-1)
+        edge_score = F.dropout(edge_score, p=self.dropout,
+                               training=self.training)
+
+        if self.edge_score_method == 'tanh':
+            edge_score = edge_score.tanh()
+        elif self.edge_score_method == 'sigmoid':
+            edge_score = edge_score.sigmoid()
+        else:
+            assert self.edge_score_method == 'log_softmax'
+            edge_score = F.log_softmax(edge_score, dim=0)
+
+        return self._merge_edges(x, edge_index, batch, edge_score)
+
+    def _merge_edges(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        batch: Tensor,
+        edge_score: Tensor,
+    ) -> Tuple[Tensor, Tensor, Tensor, UnpoolInfo]:
+
+        from scipy.sparse.csgraph import connected_components
+
+        edge_contract = edge_index[:, edge_score > self.threshhold]
+
+        adj = to_scipy_sparse_matrix(edge_contract, num_nodes=x.size(0))
+        _, cluster_np = connected_components(adj, directed=True,
+                                             connection="weak")
+
+        cluster = torch.tensor(cluster_np, dtype=torch.long, device=x.device)
+        C = one_hot(cluster)
+        A = to_dense_adj(edge_index, max_num_nodes=x.size(0)).squeeze(0)
+        S = to_dense_adj(edge_index, edge_attr=edge_score,
+                         max_num_nodes=x.size(0)).squeeze(0)
+
+        A_contract = to_dense_adj(edge_contract,
+                                  max_num_nodes=x.size(0)).squeeze(0)
+        nodes_single = ((A_contract.sum(dim=-1) +
+                         A_contract.sum(dim=-2)) == 0).nonzero()
+        S[nodes_single, nodes_single] = 1.0
+
+        x_out = (S @ C).t() @ x
+        edge_index_out, _ = dense_to_sparse((C.T @ A @ C).fill_diagonal_(0))
+        batch_out = batch.new_empty(x_out.size(0)).scatter_(0, cluster, batch)
+        unpool_info = UnpoolInfo(edge_index, cluster, batch)
+
+        return x_out, edge_index_out, batch_out, unpool_info
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.in_channels})'

torch_geometric/nn/pool/connect/base.py

@@ -66,7 +66,6 @@ class Connect(torch.nn.Module):
     """
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
-        pass
 
     def forward(
         self,

torch_geometric/nn/pool/edge_pool.py

@@ -58,7 +58,7 @@ class EdgePooling(torch.nn.Module):
         self,
         in_channels: int,
         edge_score_method: Optional[Callable] = None,
-        dropout: Optional[float] = 0.0,
+        dropout: float = 0.0,
         add_to_edge_score: float = 0.5,
     ):
         super().__init__()

torch_geometric/nn/pool/graclus.py

@@ -2,9 +2,11 @@ from typing import Optional
 
 from torch import Tensor
 
-try:
+import torch_geometric.typing
+
+if torch_geometric.typing.WITH_TORCH_CLUSTER:
     from torch_cluster import graclus_cluster
-except ImportError:
+else:
     graclus_cluster = None
 
 

torch_geometric/nn/pool/pool.py

@@ -5,12 +5,18 @@ import torch
 from torch_geometric.utils import coalesce, remove_self_loops, scatter
 
 
-def pool_edge(cluster, edge_index, edge_attr: Optional[torch.Tensor] = None):
+def pool_edge(
+    cluster,
+    edge_index,
+    edge_attr: Optional[torch.Tensor] = None,
+    reduce: Optional[str] = 'sum',
+):
     num_nodes = cluster.size(0)
     edge_index = cluster[edge_index.view(-1)].view(2, -1)
     edge_index, edge_attr = remove_self_loops(edge_index, edge_attr)
     if edge_index.numel() > 0:
-        edge_index, edge_attr = coalesce(edge_index, edge_attr, num_nodes)
+        edge_index, edge_attr = coalesce(edge_index, edge_attr, num_nodes,
+                                         reduce=reduce)
     return edge_index, edge_attr
 
 

torch_geometric/nn/pool/select/base.py

@@ -80,7 +80,6 @@ class Select(torch.nn.Module):
     """
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
-        pass
 
     def forward(self, *args, **kwargs) -> SelectOutput:
         raise NotImplementedError

torch_geometric/nn/pool/voxel_grid.py

@@ -3,11 +3,12 @@ from typing import List, Optional, Union
 import torch
 from torch import Tensor
 
+import torch_geometric.typing
 from torch_geometric.utils.repeat import repeat
 
-try:
+if torch_geometric.typing.WITH_TORCH_CLUSTER:
     from torch_cluster import grid_cluster
-except ImportError:
+else:
     grid_cluster = None
 
 

torch_geometric/nn/resolver.py

@@ -166,5 +166,5 @@ def lr_scheduler_resolver(
                 return obj
             return cls
 
-    choices = set(cls.__name__ for cls in classes)
+    choices = {cls.__name__ for cls in classes}
     raise ValueError(f"Could not resolve '{query}' among choices {choices}")

torch_geometric/nn/sequential.jinja

@@ -1,35 +1,22 @@
 import typing
-from typing import *
 
 import torch
 from torch import Tensor
 
 import torch_geometric.typing
-from torch_geometric.typing import *
+{% for module in modules %}
+from {{module}} import *
+{%- endfor %}
 
 
-class Sequential(torch.nn.Module):
-    def reset_parameters(self) -> None:
-{%- for child in children %}
-        if hasattr(self.{{child.name}}, 'reset_parameters'):
-            self.{{child.name}}.reset_parameters()
+def forward(
+    self,
+{%- for param in signature.param_dict.values() %}
+    {{param.name}}: {{param.type_repr}},
 {%- endfor %}
+) -> {{signature.return_type_repr}}:
 
-    def forward(self, {{ input_types|join(', ') }}) -> {{return_type}}:
 {%- for child in children %}
-        {{child.return_names|join(', ')}} = self.{{child.name}}({{child.param_names|join(', ')}})
+    {{child.return_names|join(', ')}} = self.{{child.name}}({{child.param_names|join(', ')}})
 {%- endfor %}
-        return {{children[-1].return_names|join(', ')}}
-
-    def __getitem__(self, idx: int) -> torch.nn.Module:
-        return getattr(self, self._module_names[idx])
-
-    def __len__(self) -> int:
-        return {{children|length}}
-
-    def __repr__(self) -> str:
-        module_reprs = [
-            f'  ({i}) - {self[i]}: {self._module_descs[i]}'
-            for i in range(len(self))
-        ]
-        return 'Sequential(\n{}\n)'.format('\n'.join(module_reprs))
+    return {{children[-1].return_names|join(', ')}}