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

torch_geometric/nn/sequential.py

@@ -1,25 +1,33 @@
+import copy
+import inspect
 import os.path as osp
 import random
-from typing import Callable, List, NamedTuple, Optional, Tuple, Union
+import sys
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    NamedTuple,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import torch
 from torch import Tensor
 
-from torch_geometric.inspector import split, type_repr
+from torch_geometric.inspector import Parameter, Signature, eval_type, split
 from torch_geometric.template import module_from_template
 
 
 class Child(NamedTuple):
     name: str
-    module: Callable
     param_names: List[str]
     return_names: List[str]
 
 
-def Sequential(
-    input_args: str,
-    modules: List[Union[Tuple[Callable, str], Callable]],
-) -> torch.nn.Module:
+class Sequential(torch.nn.Module):
     r"""An extension of the :class:`torch.nn.Sequential` container in order to
     define a sequential GNN model.
 
@@ -69,79 +77,197 @@ def Sequential(
 
     Args:
         input_args (str): The input arguments of the model.
-        modules ([(str, Callable) or Callable]): A list of modules (with
+        modules ([(Callable, str) or Callable]): A list of modules (with
             optional function header definitions). Alternatively, an
             :obj:`OrderedDict` of modules (and function header definitions) can
             be passed.
     """
-    signature = input_args.split('->')
-    if len(signature) == 1:
-        input_args = signature[0]
-        return_type = type_repr(Tensor, globals())
-    elif len(signature) == 2:
-        input_args, return_type = signature[0], signature[1].strip()
-    else:
-        raise ValueError(f"Failed to parse arguments (got '{input_args}')")
-
-    input_types = split(input_args, sep=',')
-    if len(input_types) == 0:
-        raise ValueError(f"Failed to parse arguments (got '{input_args}')")
-
-    if not isinstance(modules, dict):
-        modules = {f'module_{i}': module for i, module in enumerate(modules)}
-    if len(modules) == 0:
-        raise ValueError("'Sequential' expected a non-empty list of modules")
-
-    children: List[Child] = []
-    for i, (name, module) in enumerate(modules.items()):
-        desc: Optional[str] = None
-        if isinstance(module, (tuple, list)):
-            if len(module) == 1:
-                module = module[0]
-            elif len(module) == 2:
-                module, desc = module
-            else:
-                raise ValueError(f"Expected tuple of length 2 (got {module})")
-
-        if i == 0 and desc is None:
-            raise ValueError("Requires signature for first module")
-        if not callable(module):
-            raise ValueError(f"Expected callable module (got {module})")
-        if desc is not None and not isinstance(desc, str):
-            raise ValueError(f"Expected type hint representation (got {desc})")
-
-        if desc is not None:
-            signature = desc.split('->')
-            if len(signature) != 2:
-                raise ValueError(f"Failed to parse arguments (got '{desc}')")
-            param_names = [v.strip() for v in signature[0].split(',')]
-            return_names = [v.strip() for v in signature[1].split(',')]
-            child = Child(name, module, param_names, return_names)
+    _children: List[Child]
+
+    def __init__(
+        self,
+        input_args: str,
+        modules: List[Union[Tuple[Callable, str], Callable]],
+    ) -> None:
+        super().__init__()
+
+        caller_path = inspect.stack()[1].filename
+        self._caller_module = osp.splitext(osp.basename(caller_path))[0]
+
+        _globals = copy.copy(globals())
+        _globals.update(sys.modules['__main__'].__dict__)
+        if self._caller_module in sys.modules:
+            _globals.update(sys.modules[self._caller_module].__dict__)
+
+        signature = input_args.split('->')
+        if len(signature) == 1:
+            args_repr = signature[0]
+            return_type_repr = 'Tensor'
+            return_type = Tensor
+        elif len(signature) == 2:
+            args_repr = signature[0]
+            return_type_repr = signature[1].strip()
+            return_type = eval_type(return_type_repr, _globals)
         else:
-            param_names = children[-1].return_names
-            child = Child(name, module, param_names, param_names)
-
-        children.append(child)
-
-    uid = '%06x' % random.randrange(16**6)
-    root_dir = osp.dirname(osp.realpath(__file__))
-    module = module_from_template(
-        module_name=f'torch_geometric.nn.sequential_{uid}',
-        template_path=osp.join(root_dir, 'sequential.jinja'),
-        tmp_dirname='sequential',
-        # Keyword arguments:
-        input_types=input_types,
-        return_type=return_type,
-        children=children,
-    )
-
-    model = module.Sequential()
-    model._module_names = [child.name for child in children]
-    model._module_descs = [
-        f"{', '.join(child.param_names)} -> {', '.join(child.return_names)}"
-        for child in children
-    ]
-    for child in children:
-        setattr(model, child.name, child.module)
-
-    return model
+            raise ValueError(f"Failed to parse arguments (got '{input_args}')")
+
+        param_dict: Dict[str, Parameter] = {}
+        for arg in split(args_repr, sep=','):
+            signature = arg.split(':')
+            if len(signature) == 1:
+                name = signature[0].strip()
+                param_dict[name] = Parameter(
+                    name=name,
+                    type=Tensor,
+                    type_repr='Tensor',
+                    default=inspect._empty,
+                )
+            elif len(signature) == 2:
+                name = signature[0].strip()
+                param_dict[name] = Parameter(
+                    name=name,
+                    type=eval_type(signature[1].strip(), _globals),
+                    type_repr=signature[1].strip(),
+                    default=inspect._empty,
+                )
+            else:
+                raise ValueError(f"Failed to parse argument "
+                                 f"(got '{arg.strip()}')")
+
+        self.signature = Signature(param_dict, return_type, return_type_repr)
+
+        if not isinstance(modules, dict):
+            modules = {
+                f'module_{i}': module
+                for i, module in enumerate(modules)
+            }
+        if len(modules) == 0:
+            raise ValueError(f"'{self.__class__.__name__}' expects a "
+                             f"non-empty list of modules")
+
+        self._children: List[Child] = []
+        for i, (name, module) in enumerate(modules.items()):
+            desc: Optional[str] = None
+            if isinstance(module, (tuple, list)):
+                if len(module) == 1:
+                    module = module[0]
+                elif len(module) == 2:
+                    module, desc = module
+                else:
+                    raise ValueError(f"Expected tuple of length 2 "
+                                     f"(got {module})")
+
+            if i == 0 and desc is None:
+                raise ValueError("Signature for first module required")
+            if not callable(module):
+                raise ValueError(f"Expected callable module (got {module})")
+            if desc is not None and not isinstance(desc, str):
+                raise ValueError(f"Expected type hint representation "
+                                 f"(got {desc})")
+
+            if desc is not None:
+                signature = desc.split('->')
+                if len(signature) != 2:
+                    raise ValueError(
+                        f"Failed to parse arguments (got '{desc}')")
+                param_names = [v.strip() for v in signature[0].split(',')]
+                return_names = [v.strip() for v in signature[1].split(',')]
+                child = Child(name, param_names, return_names)
+            else:
+                param_names = self._children[-1].return_names
+                child = Child(name, param_names, param_names)
+
+            setattr(self, name, module)
+            self._children.append(child)
+
+        self._set_jittable_template()
+
+    def reset_parameters(self) -> None:
+        r"""Resets all learnable parameters of the module."""
+        for child in self._children:
+            module = getattr(self, child.name)
+            if hasattr(module, 'reset_parameters'):
+                module.reset_parameters()
+
+    def __len__(self) -> int:
+        return len(self._children)
+
+    def __getitem__(self, idx: int) -> torch.nn.Module:
+        return getattr(self, self._children[idx].name)
+
+    def __setstate__(self, data: Dict[str, Any]) -> None:
+        super().__setstate__(data)
+        self._set_jittable_template()
+
+    def __repr__(self) -> str:
+        module_descs = [
+            f"{', '.join(c.param_names)} -> {', '.join(c.return_names)}"
+            for c in self._children
+        ]
+        module_reprs = [
+            f'  ({i}) - {self[i]}: {module_descs[i]}' for i in range(len(self))
+        ]
+        return '{}(\n{}\n)'.format(
+            self.__class__.__name__,
+            '\n'.join(module_reprs),
+        )
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
+        """"""  # noqa: D419
+        value_dict = {
+            name: arg
+            for name, arg in zip(self.signature.param_dict.keys(), args)
+        }
+        for key, arg in kwargs.items():
+            if key in value_dict:
+                raise TypeError(f"'{self.__class__.__name__}' got multiple "
+                                f"values for argument '{key}'")
+            value_dict[key] = arg
+
+        for child in self._children:
+            args = [value_dict[name] for name in child.param_names]
+            outs = getattr(self, child.name)(*args)
+            if len(child.return_names) == 1:
+                value_dict[child.return_names[0]] = outs
+            else:
+                for name, out in zip(child.return_names, outs):
+                    value_dict[name] = out
+
+        return outs
+
+    # TorchScript Support #####################################################
+
+    def _set_jittable_template(self, raise_on_error: bool = False) -> None:
+        try:  # Optimize `forward()` via `*.jinja` templates:
+            if ('forward' in self.__class__.__dict__ and
+                    self.__class__.__dict__['forward'] != Sequential.forward):
+                raise ValueError("Cannot compile custom 'forward' method")
+
+            root_dir = osp.dirname(osp.realpath(__file__))
+            uid = '%06x' % random.randrange(16**6)
+            jinja_prefix = f'{self.__module__}_{self.__class__.__name__}_{uid}'
+            module = module_from_template(
+                module_name=jinja_prefix,
+                template_path=osp.join(root_dir, 'sequential.jinja'),
+                tmp_dirname='sequential',
+                # Keyword arguments:
+                modules=[self._caller_module],
+                signature=self.signature,
+                children=self._children,
+            )
+
+            self.forward = module.forward.__get__(self)
+
+            # NOTE We override `forward` on the class level here in order to
+            # support `torch.jit.trace` - this is generally dangerous to do,
+            # and limits `torch.jit.trace` to a single `Sequential` module:
+            self.__class__.forward = module.forward
+        except Exception as e:  # pragma: no cover
+            if raise_on_error:
+                raise e
+
+    def __prepare_scriptable__(self) -> 'Sequential':
+        # Prevent type sharing when scripting `Sequential` modules:
+        type_store = torch.jit._recursive.concrete_type_store.type_store
+        type_store.pop(self.__class__, None)
+        return self

torch_geometric/nn/summary.py

@@ -141,7 +141,7 @@ def get_shape(inputs: Any) -> str:
 def postprocess(info_list: List[dict]) -> List[dict]:
     for idx, info in enumerate(info_list):
         depth = info['depth']
-        if idx > 0:  # root module (0) is exclued
+        if idx > 0:  # root module (0) is excluded
             if depth == 1:
                 prefix = '├─'
             else:

torch_geometric/nn/to_hetero_with_bases_transformer.py

@@ -272,7 +272,6 @@ class ToHeteroWithBasesTransformer(Transformer):
                     args=(value, self.find_by_name('edge_offset_dict')),
                     name=f'{value.name}__split')
 
-                pass
             elif isinstance(value, Node):
                 self.graph.inserting_before(node)
                 return self.graph.create_node(
@@ -309,6 +308,24 @@ class ToHeteroWithBasesTransformer(Transformer):
 ###############################################################################
 
 
+# We make use of a post-message computation hook to inject the
+# basis re-weighting for each individual edge type.
+# This currently requires us to set `conv.fuse = False`, which leads
+# to a materialization of messages.
+def hook(module, inputs, output):
+    assert isinstance(module._edge_type, Tensor)
+    if module._edge_type.size(0) != output.size(-2):
+        raise ValueError(
+            f"Number of messages ({output.size(0)}) does not match "
+            f"with the number of original edges "
+            f"({module._edge_type.size(0)}). Does your message "
+            f"passing layer create additional self-loops? Try to "
+            f"remove them via 'add_self_loops=False'")
+    weight = module.edge_type_weight.view(-1)[module._edge_type]
+    weight = weight.view([1] * (output.dim() - 2) + [-1, 1])
+    return weight * output
+
+
 class HeteroBasisConv(torch.nn.Module):
     # A wrapper layer that applies the basis-decomposition technique to a
     # heterogeneous graph.
@@ -319,23 +336,6 @@ class HeteroBasisConv(torch.nn.Module):
         self.num_relations = num_relations
         self.num_bases = num_bases
 
-        # We make use of a post-message computation hook to inject the
-        # basis re-weighting for each individual edge type.
-        # This currently requires us to set `conv.fuse = False`, which leads
-        # to a materialization of messages.
-        def hook(module, inputs, output):
-            assert isinstance(module._edge_type, Tensor)
-            if module._edge_type.size(0) != output.size(-2):
-                raise ValueError(
-                    f"Number of messages ({output.size(0)}) does not match "
-                    f"with the number of original edges "
-                    f"({module._edge_type.size(0)}). Does your message "
-                    f"passing layer create additional self-loops? Try to "
-                    f"remove them via 'add_self_loops=False'")
-            weight = module.edge_type_weight.view(-1)[module._edge_type]
-            weight = weight.view([1] * (output.dim() - 2) + [-1, 1])
-            return weight * output
-
         params = list(module.parameters())
         device = params[0].device if len(params) > 0 else 'cpu'
 
@@ -468,7 +468,7 @@ def get_edge_type(
 ###############################################################################
 
 # These methods are used to group the individual type-wise components into a
-# unfied single representation.
+# unified single representation.
 
 
 def group_node_placeholder(input_dict: Dict[NodeType, Tensor],

torch_geometric/profile/__init__.py

@@ -20,6 +20,7 @@ from .utils import (
     get_gpu_memory_from_nvidia_smi,
     get_model_size,
 )
+from .nvtx import nvtxit
 
 __all__ = [
     'profileit',
@@ -38,6 +39,7 @@ __all__ = [
     'get_gpu_memory_from_nvidia_smi',
     'get_gpu_memory_from_ipex',
     'benchmark',
+    'nvtxit',
 ]
 
 classes = __all__

torch_geometric/profile/nvtx.py

@@ -0,0 +1,66 @@
+from functools import wraps
+from typing import Optional
+
+import torch
+
+CUDA_PROFILE_STARTED = False
+
+
+def begin_cuda_profile():
+    global CUDA_PROFILE_STARTED
+    prev_state = CUDA_PROFILE_STARTED
+    if prev_state is False:
+        CUDA_PROFILE_STARTED = True
+        torch.cuda.cudart().cudaProfilerStart()
+    return prev_state
+
+
+def end_cuda_profile(prev_state: bool):
+    global CUDA_PROFILE_STARTED
+    CUDA_PROFILE_STARTED = prev_state
+    if prev_state is False:
+        torch.cuda.cudart().cudaProfilerStop()
+
+
+def nvtxit(name: Optional[str] = None, n_warmups: int = 0,
+           n_iters: Optional[int] = None):
+    """Enables NVTX profiling for a function.
+
+    Args:
+        name (Optional[str], optional): Name to give the reference frame for
+            the function being wrapped. Defaults to the name of the
+            function in code.
+        n_warmups (int, optional): Number of iters to call that function
+            before starting. Defaults to 0.
+        n_iters (Optional[int], optional): Number of iters of that function to
+            record. Defaults to all of them.
+    """
+    def nvtx(func):
+
+        nonlocal name
+        iters_so_far = 0
+        if name is None:
+            name = func.__name__
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            nonlocal iters_so_far
+            if not torch.cuda.is_available():
+                return func(*args, **kwargs)
+            elif iters_so_far < n_warmups:
+                iters_so_far += 1
+                return func(*args, **kwargs)
+            elif n_iters is None or iters_so_far < n_iters + n_warmups:
+                prev_state = begin_cuda_profile()
+                torch.cuda.nvtx.range_push(f"{name}_{iters_so_far}")
+                result = func(*args, **kwargs)
+                torch.cuda.nvtx.range_pop()
+                end_cuda_profile(prev_state)
+                iters_so_far += 1
+                return result
+            else:
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return nvtx

torch_geometric/profile/profiler.py

@@ -5,6 +5,8 @@ from typing import Any, List, NamedTuple, Optional, Tuple
 import torch
 import torch.profiler as torch_profiler
 
+import torch_geometric.typing
+
 # predefined namedtuple for variable setting (global template)
 Trace = namedtuple('Trace', ['path', 'leaf', 'module'])
 
@@ -325,41 +327,50 @@ def _flatten_tree(t, depth=0):
 
 
 def _build_measure_tuple(events: List, occurrences: List) -> NamedTuple:
+    device_str = 'device' if torch_geometric.typing.WITH_PT24 else 'cuda'
+
     # memory profiling supported in torch >= 1.6
     self_cpu_memory = None
     has_self_cpu_memory = any(
         hasattr(e, "self_cpu_memory_usage") for e in events)
     if has_self_cpu_memory:
         self_cpu_memory = sum(
-            [getattr(e, "self_cpu_memory_usage", 0) for e in events])
+            [getattr(e, "self_cpu_memory_usage", 0) or 0 for e in events])
     cpu_memory = None
     has_cpu_memory = any(hasattr(e, "cpu_memory_usage") for e in events)
     if has_cpu_memory:
-        cpu_memory = sum([getattr(e, "cpu_memory_usage", 0) for e in events])
+        cpu_memory = sum(
+            [getattr(e, "cpu_memory_usage", 0) or 0 for e in events])
     self_cuda_memory = None
     has_self_cuda_memory = any(
-        hasattr(e, "self_cuda_memory_usage") for e in events)
+        hasattr(e, f"self_{device_str}_memory_usage") for e in events)
     if has_self_cuda_memory:
-        self_cuda_memory = sum(
-            [getattr(e, "self_cuda_memory_usage", 0) for e in events])
+        self_cuda_memory = sum([
+            getattr(e, f"self_{device_str}_memory_usage", 0) or 0
+            for e in events
+        ])
     cuda_memory = None
-    has_cuda_memory = any(hasattr(e, "cuda_memory_usage") for e in events)
+    has_cuda_memory = any(
+        hasattr(e, f"{device_str}_memory_usage") for e in events)
     if has_cuda_memory:
-        cuda_memory = sum([getattr(e, "cuda_memory_usage", 0) for e in events])
+        cuda_memory = sum(
+            [getattr(e, f"{device_str}_memory_usage", 0) or 0 for e in events])
 
     # self CUDA time supported in torch >= 1.7
     self_cuda_total = None
     has_self_cuda_time = any(
-        hasattr(e, "self_cuda_time_total") for e in events)
+        hasattr(e, f"self_{device_str}_time_total") for e in events)
     if has_self_cuda_time:
-        self_cuda_total = sum(
-            [getattr(e, "self_cuda_time_total", 0) for e in events])
+        self_cuda_total = sum([
+            getattr(e, f"self_{device_str}_time_total", 0) or 0 for e in events
+        ])
 
     return Measure(
-        self_cpu_total=sum([e.self_cpu_time_total for e in events]),
-        cpu_total=sum([e.cpu_time_total for e in events]),
+        self_cpu_total=sum([e.self_cpu_time_total or 0 for e in events]),
+        cpu_total=sum([e.cpu_time_total or 0 for e in events]),
         self_cuda_total=self_cuda_total,
-        cuda_total=sum([e.cuda_time_total for e in events]),
+        cuda_total=sum(
+            [getattr(e, f"{device_str}_time_total") or 0 for e in events]),
         self_cpu_memory=self_cpu_memory,
         cpu_memory=cpu_memory,
         self_cuda_memory=self_cuda_memory,
@@ -434,10 +445,10 @@ def format_time(time_us: int) -> str:
     US_IN_SECOND = 1000.0 * 1000.0
     US_IN_MS = 1000.0
     if time_us >= US_IN_SECOND:
-        return '{:.3f}s'.format(time_us / US_IN_SECOND)
+        return f'{time_us / US_IN_SECOND:.3f}s'
     if time_us >= US_IN_MS:
-        return '{:.3f}ms'.format(time_us / US_IN_MS)
-    return '{:.3f}us'.format(time_us)
+        return f'{time_us / US_IN_MS:.3f}ms'
+    return f'{time_us:.3f}us'
 
 
 def format_memory(nbytes: int) -> str:
@@ -446,10 +457,10 @@ def format_memory(nbytes: int) -> str:
     MB = 1024 * KB
     GB = 1024 * MB
     if (abs(nbytes) >= GB):
-        return '{:.2f} Gb'.format(nbytes * 1.0 / GB)
+        return f'{nbytes * 1.0 / GB:.2f} Gb'
     elif (abs(nbytes) >= MB):
-        return '{:.2f} Mb'.format(nbytes * 1.0 / MB)
+        return f'{nbytes * 1.0 / MB:.2f} Mb'
     elif (abs(nbytes) >= KB):
-        return '{:.2f} Kb'.format(nbytes * 1.0 / KB)
+        return f'{nbytes * 1.0 / KB:.2f} Kb'
     else:
         return str(nbytes) + ' b'

torch_geometric/resolver.py

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