Source code for torch_geometric.transforms.pad

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
import torch.nn.functional as F

from torch_geometric.data import Data, HeteroData
from torch_geometric.data.datapipes import functional_transform
from torch_geometric.data.storage import EdgeStorage, NodeStorage
from torch_geometric.transforms import BaseTransform
from torch_geometric.typing import EdgeType, NodeType


class Padding(ABC):
    r"""An abstract class for specifying padding values."""
    @abstractmethod
    def get_value(
        self,
        store_type: Optional[Union[NodeType, EdgeType]] = None,
        attr_name: Optional[str] = None,
    ) -> Union[int, float]:
        pass


@dataclass(init=False)
class UniformPadding(Padding):
    r"""Uniform padding independent of attribute name or node/edge type.

    Args:
        value (int or float, optional): The value to be used for padding.
            (default: :obj:`0.0`)
    """
    value: Union[int, float] = 0.0

    def __init__(self, value: Union[int, float] = 0.0):
        self.value = value

        if not isinstance(self.value, (int, float)):
            raise ValueError(f"Expected 'value' to be an integer or float "
                             f"(got '{type(value)}'")

    def get_value(
        self,
        store_type: Optional[Union[NodeType, EdgeType]] = None,
        attr_name: Optional[str] = None,
    ) -> Union[int, float]:
        return self.value


@dataclass(init=False)
class MappingPadding(Padding):
    r"""An abstract class for specifying different padding values."""
    values: Dict[Any, Padding]
    default: UniformPadding

    def __init__(
        self,
        values: Dict[Any, Union[int, float, Padding]],
        default: Union[int, float] = 0.0,
    ):
        if not isinstance(values, dict):
            raise ValueError(f"Expected 'values' to be a dictionary "
                             f"(got '{type(values)}'")

        self.values = {
            key: UniformPadding(val) if isinstance(val, (int, float)) else val
            for key, val in values.items()
        }
        self.default = UniformPadding(default)

        for key, value in self.values.items():
            self.validate_key_value(key, value)

    def validate_key_value(self, key: Any, value: Any) -> None:
        pass


class AttrNamePadding(MappingPadding):
    r"""Padding dependent on attribute names.

    Args:
        values (dict): The mapping from attribute names to padding values.
        default (int or float, optional): The padding value to use for
            attribute names not specified in :obj:`values`.
            (default: :obj:`0.0`)
    """
    def validate_key_value(self, key: Any, value: Any) -> None:
        if not isinstance(key, str):
            raise ValueError(f"Expected the attribute name '{key}' to be a "
                             f"string (got '{type(key)}')")

        if not isinstance(value, UniformPadding):
            raise ValueError(f"Expected the value of '{key}' to be of "
                             f"type 'UniformPadding' (got '{type(value)}')")

    def get_value(
        self,
        store_type: Optional[Union[NodeType, EdgeType]] = None,
        attr_name: Optional[str] = None,
    ) -> Union[int, float]:
        padding = self.values.get(attr_name, self.default)
        return padding.get_value()


class NodeTypePadding(MappingPadding):
    r"""Padding dependent on node types.

    Args:
        values (dict): The mapping from node types to padding values.
        default (int or float, optional): The padding value to use for node
            types not specified in :obj:`values`. (default: :obj:`0.0`)
    """
    def validate_key_value(self, key: Any, value: Any) -> None:
        if not isinstance(key, str):
            raise ValueError(f"Expected the node type '{key}' to be a string "
                             f"(got '{type(key)}')")

        if not isinstance(value, (UniformPadding, AttrNamePadding)):
            raise ValueError(f"Expected the value of '{key}' to be of "
                             f"type 'UniformPadding' or 'AttrNamePadding' "
                             f"(got '{type(value)}')")

    def get_value(
        self,
        store_type: Optional[Union[NodeType, EdgeType]] = None,
        attr_name: Optional[str] = None,
    ) -> Union[int, float]:
        padding = self.values.get(store_type, self.default)
        return padding.get_value(attr_name=attr_name)


class EdgeTypePadding(MappingPadding):
    r"""Padding dependent on node types.

    Args:
        values (dict): The mapping from edge types to padding values.
        default (int or float, optional): The padding value to use for edge
            types not specified in :obj:`values`. (default: :obj:`0.0`)
    """
    def validate_key_value(self, key: Any, value: Any) -> None:
        if not isinstance(key, tuple):
            raise ValueError(f"Expected the edge type '{key}' to be a tuple "
                             f"(got '{type(key)}')")

        if len(key) != 3:
            raise ValueError(f"Expected the edge type '{key}' to hold exactly "
                             f"three elements (got {len(key)})")

        if not isinstance(value, (UniformPadding, AttrNamePadding)):
            raise ValueError(f"Expected the value of '{key}' to be of "
                             f"type 'UniformPadding' or 'AttrNamePadding' "
                             f"(got '{type(value)}')")

    def get_value(
        self,
        store_type: Optional[Union[NodeType, EdgeType]] = None,
        attr_name: Optional[str] = None,
    ) -> Union[int, float]:
        padding = self.values.get(store_type, self.default)
        return padding.get_value(attr_name=attr_name)


class _NumNodes:
    def __init__(
        self,
        value: Union[int, Dict[NodeType, int], None],
    ) -> None:
        self.value = value

    def get_value(self, key: Optional[NodeType] = None) -> Optional[int]:
        if self.value is None or isinstance(self.value, int):
            return self.value
        assert isinstance(key, str)
        return self.value[key]


class _NumEdges:
    def __init__(
        self,
        value: Union[int, Dict[EdgeType, int], None],
        num_nodes: _NumNodes,
    ) -> None:

        if value is None:
            if isinstance(num_nodes.value, int):
                value = num_nodes.value * num_nodes.value
            else:
                value = {}

        self.value = value
        self.num_nodes = num_nodes

    def get_value(self, key: Optional[EdgeType] = None) -> Optional[int]:
        if self.value is None or isinstance(self.value, int):
            return self.value

        assert isinstance(key, tuple) and len(key) == 3
        if key not in self.value:
            num_src_nodes = self.num_nodes.get_value(key[0])
            num_dst_nodes = self.num_nodes.get_value(key[-1])
            assert num_src_nodes is not None and num_dst_nodes is not None
            self.value[key] = num_src_nodes * num_dst_nodes

        return self.value[key]


[docs]@functional_transform('pad') class Pad(BaseTransform): r"""Applies padding to enforce consistent tensor shapes (functional name: :obj:`pad`). This transform will pad node and edge features up to a maximum allowed size in the node or edge feature dimension. By default :obj:`0.0` is used as the padding value and can be configured by setting :obj:`node_pad_value` and :obj:`edge_pad_value`. In case of applying :class:`Pad` to a :class:`~torch_geometric.data.Data` object, the :obj:`node_pad_value` value (or :obj:`edge_pad_value`) can be either: * an int, float or object of :class:`UniformPadding` class for cases when all attributes are going to be padded with the same value; * an object of :class:`AttrNamePadding` class for cases when padding is going to differ based on attribute names. In case of applying :class:`Pad` to a :class:`~torch_geometric.data.HeteroData` object, the :obj:`node_pad_value` value (or :obj:`edge_pad_value`) can be either: * an int, float or object of :class:`UniformPadding` class for cases when all attributes of all node (or edge) stores are going to be padded with the same value; * an object of :class:`AttrNamePadding` class for cases when padding is going to differ based on attribute names (but not based on node or edge types); * an object of class :class:`NodeTypePadding` or :class:`EdgeTypePadding` for cases when padding values are going to differ based on node or edge types. Padding values can also differ based on attribute names for a given node or edge type by using :class:`AttrNamePadding` objects as values of its `values` argument. Note that in order to allow for consistent padding across all graphs in a dataset, below conditions must be met: * if :obj:`max_num_nodes` is a single value, it must be greater than or equal to the maximum number of nodes of any graph in the dataset; * if :obj:`max_num_nodes` is a dictionary, value for every node type must be greater than or equal to the maximum number of this type nodes of any graph in the dataset. Example below shows how to create a :class:`Pad` transform for an :class:`~torch_geometric.data.HeteroData` object. The object is padded to have :obj:`10` nodes of type :obj:`v0`, :obj:`20` nodes of type :obj:`v1` and :obj:`30` nodes of type :obj:`v2`. It is padded to have :obj:`80` edges of type :obj:`('v0', 'e0', 'v1')`. All the attributes of the :obj:`v0` nodes are padded using a value of :obj:`3.0`. The :obj:`x` attribute of the :obj:`v1` node type is padded using a value of :obj:`-1.0`, and the other attributes of this node type are padded using a value of :obj:`0.5`. All the attributes of node types other than :obj:`v0` and :obj:`v1` are padded using a value of :obj:`1.0`. All the attributes of the :obj:`('v0', 'e0', 'v1')` edge type are padded using a value of :obj:`3.5`. The :obj:`edge_attr` attributes of the :obj:`('v1', 'e0', 'v0')` edge type are padded using a value of :obj:`-1.5`, and any other attributes of this edge type are padded using a value of :obj:`5.5`. All the attributes of edge types other than these two are padded using a value of :obj:`1.5`. .. code-block:: python num_nodes = {'v0': 10, 'v1': 20, 'v2':30} num_edges = {('v0', 'e0', 'v1'): 80} node_padding = NodeTypePadding({ 'v0': 3.0, 'v1': AttrNamePadding({'x': -1.0}, default=0.5), }, default=1.0) edge_padding = EdgeTypePadding({ ('v0', 'e0', 'v1'): 3.5, ('v1', 'e0', 'v0'): AttrNamePadding({'edge_attr': -1.5}, default=5.5), }, default=1.5) transform = Pad(num_nodes, num_edges, node_padding, edge_padding) Args: max_num_nodes (int or dict): The number of nodes after padding. In heterogeneous graphs, may also take in a dictionary denoting the number of nodes for specific node types. max_num_edges (int or dict, optional): The number of edges after padding. In heterogeneous graphs, may also take in a dictionary denoting the number of edges for specific edge types. (default: :obj:`None`) node_pad_value (int or float or Padding, optional): The fill value to use for node features. (default: :obj:`0.0`) edge_pad_value (int or float or Padding, optional): The fill value to use for edge features. (default: :obj:`0.0`) The :obj:`edge_index` tensor is padded with with the index of the first padded node (which represents a set of self-loops on the padded node). (default: :obj:`0.0`) mask_pad_value (bool, optional): The fill value to use for :obj:`train_mask`, :obj:`val_mask` and :obj:`test_mask` attributes (default: :obj:`False`). add_pad_mask (bool, optional): If set to :obj:`True`, will attach node-level :obj:`pad_node_mask` and edge-level :obj:`pad_edge_mask` attributes to the output which indicates which elements in the data are real (represented by :obj:`True`) and which were added as a result of padding (represented by :obj:`False`). (default: :obj:`False`) exclude_keys ([str], optional): Keys to be removed from the input data object. (default: :obj:`None`) """ def __init__( self, max_num_nodes: Union[int, Dict[NodeType, int]], max_num_edges: Optional[Union[int, Dict[EdgeType, int]]] = None, node_pad_value: Union[int, float, Padding] = 0.0, edge_pad_value: Union[int, float, Padding] = 0.0, mask_pad_value: bool = False, add_pad_mask: bool = False, exclude_keys: Optional[List[str]] = None, ): self.max_num_nodes = _NumNodes(max_num_nodes) self.max_num_edges = _NumEdges(max_num_edges, self.max_num_nodes) self.node_pad: Padding if not isinstance(node_pad_value, Padding): self.node_pad = UniformPadding(node_pad_value) else: self.node_pad = node_pad_value self.edge_pad: Padding if not isinstance(edge_pad_value, Padding): self.edge_pad = UniformPadding(edge_pad_value) else: self.edge_pad = edge_pad_value self.node_additional_attrs_pad = { key: mask_pad_value for key in ['train_mask', 'val_mask', 'test_mask'] } self.add_pad_mask = add_pad_mask self.exclude_keys = set(exclude_keys or []) def __should_pad_node_attr(self, attr_name: str) -> bool: if attr_name in self.node_additional_attrs_pad: return True if self.exclude_keys is None or attr_name not in self.exclude_keys: return True return False def __should_pad_edge_attr(self, attr_name: str) -> bool: if self.max_num_edges.value is None: return False if attr_name == 'edge_index': return True if self.exclude_keys is None or attr_name not in self.exclude_keys: return True return False def __get_node_padding( self, attr_name: str, node_type: Optional[NodeType] = None, ) -> Union[int, float]: if attr_name in self.node_additional_attrs_pad: return self.node_additional_attrs_pad[attr_name] return self.node_pad.get_value(node_type, attr_name) def __get_edge_padding( self, attr_name: str, edge_type: Optional[EdgeType] = None, ) -> Union[int, float]: return self.edge_pad.get_value(edge_type, attr_name) def forward( self, data: Union[Data, HeteroData], ) -> Union[Data, HeteroData]: if isinstance(data, Data): assert isinstance(self.node_pad, (UniformPadding, AttrNamePadding)) assert isinstance(self.edge_pad, (UniformPadding, AttrNamePadding)) for key in self.exclude_keys: del data[key] num_nodes = data.num_nodes assert num_nodes is not None self.__pad_edge_store(data._store, data.__cat_dim__, num_nodes) self.__pad_node_store(data._store, data.__cat_dim__) data.num_nodes = self.max_num_nodes.get_value() else: assert isinstance( self.node_pad, (UniformPadding, AttrNamePadding, NodeTypePadding)) assert isinstance( self.edge_pad, (UniformPadding, AttrNamePadding, EdgeTypePadding)) for edge_type, edge_store in data.edge_items(): for key in self.exclude_keys: del edge_store[key] src_node_type, _, dst_node_type = edge_type num_src_nodes = data[src_node_type].num_nodes num_dst_nodes = data[dst_node_type].num_nodes assert num_src_nodes is not None and num_dst_nodes is not None self.__pad_edge_store(edge_store, data.__cat_dim__, (num_src_nodes, num_dst_nodes), edge_type) for node_type, node_store in data.node_items(): for key in self.exclude_keys: del node_store[key] self.__pad_node_store(node_store, data.__cat_dim__, node_type) data[node_type].num_nodes = self.max_num_nodes.get_value( node_type) return data def __pad_node_store( self, store: NodeStorage, get_dim_fn: Callable, node_type: Optional[NodeType] = None, ) -> None: attrs_to_pad = [key for key in store.keys() if store.is_node_attr(key)] if len(attrs_to_pad) == 0: return num_target_nodes = self.max_num_nodes.get_value(node_type) assert num_target_nodes is not None assert store.num_nodes is not None assert num_target_nodes >= store.num_nodes, \ f'The number of nodes after padding ({num_target_nodes}) cannot ' \ f'be lower than the number of nodes in the data object ' \ f'({store.num_nodes}).' num_pad_nodes = num_target_nodes - store.num_nodes if self.add_pad_mask: pad_node_mask = torch.ones(num_target_nodes, dtype=torch.bool) pad_node_mask[store.num_nodes:] = False store.pad_node_mask = pad_node_mask for attr_name in attrs_to_pad: attr = store[attr_name] pad_value = self.__get_node_padding(attr_name, node_type) dim = get_dim_fn(attr_name, attr) store[attr_name] = self._pad_tensor_dim(attr, dim, num_pad_nodes, pad_value) def __pad_edge_store( self, store: EdgeStorage, get_dim_fn: Callable, num_nodes: Union[int, Tuple[int, int]], edge_type: Optional[EdgeType] = None, ) -> None: attrs_to_pad = { attr for attr in store.keys() if store.is_edge_attr(attr) and self.__should_pad_edge_attr(attr) } if not attrs_to_pad: return num_target_edges = self.max_num_edges.get_value(edge_type) assert num_target_edges is not None assert num_target_edges >= store.num_edges, \ f'The number of edges after padding ({num_target_edges}) cannot ' \ f'be lower than the number of edges in the data object ' \ f'({store.num_edges}).' num_pad_edges = num_target_edges - store.num_edges if self.add_pad_mask: pad_edge_mask = torch.ones(num_target_edges, dtype=torch.bool) pad_edge_mask[store.num_edges:] = False store.pad_edge_mask = pad_edge_mask if isinstance(num_nodes, tuple): src_pad_value, dst_pad_value = num_nodes else: src_pad_value = dst_pad_value = num_nodes for attr_name in attrs_to_pad: attr = store[attr_name] dim = get_dim_fn(attr_name, attr) if attr_name == 'edge_index': store[attr_name] = self._pad_edge_index( attr, num_pad_edges, src_pad_value, dst_pad_value) else: pad_value = self.__get_edge_padding(attr_name, edge_type) store[attr_name] = self._pad_tensor_dim( attr, dim, num_pad_edges, pad_value) @staticmethod def _pad_tensor_dim(input: torch.Tensor, dim: int, length: int, pad_value: float) -> torch.Tensor: r"""Pads the input tensor in the specified dim with a constant value of the given length. """ pads = [0] * (2 * input.ndim) pads[-2 * dim - 1] = length return F.pad(input, pads, 'constant', pad_value) @staticmethod def _pad_edge_index(input: torch.Tensor, length: int, src_pad_value: float, dst_pad_value: float) -> torch.Tensor: r"""Pads the edges :obj:`edge_index` feature with values specified separately for src and dst nodes. """ pads = [0, length, 0, 0] padded = F.pad(input, pads, 'constant', src_pad_value) if src_pad_value != dst_pad_value: padded[1, input.shape[1]:] = dst_pad_value return padded def __repr__(self) -> str: s = f'{self.__class__.__name__}(' s += f'max_num_nodes={self.max_num_nodes.value}, ' s += f'max_num_edges={self.max_num_edges.value}, ' s += f'node_pad_value={self.node_pad}, ' s += f'edge_pad_value={self.edge_pad})' return s