from collections import defaultdict from collections.abc import Mapping from typing import Any, List, Optional, Sequence, Tuple, Union import numpy as np import torch.utils.data import torch_geometric import torch_sparse from torch.utils.data.dataloader import default_collate from torch_geometric.data import Batch, Dataset from torch_geometric.data.data import BaseData from torch_geometric.data.datapipes import DatasetAdapter from torch_geometric.data.on_disk_dataset import OnDiskDataset from torch_geometric.data.storage import BaseStorage from torch_geometric.typing import SparseTensor, TensorFrame, torch_frame from torch_geometric.utils import is_sparse, is_torch_sparse_tensor from torch_geometric.utils.sparse import cat FLOAT_PADDING_VALUE = 1e-8 NON_FLOAT_PADDING_VALUE = -1 def _dense_pad_tensor( key, values, float_padding_value: float = FLOAT_PADDING_VALUE, non_float_padding_value: int = NON_FLOAT_PADDING_VALUE, max_size: Optional[int] = None, ) -> Tuple[Any, Any]: """Densely pads a list of `torch.Tensor` along a new `batch_dim=0`. Args: key (str): The key of the attribute to be padded. values (list): A list of values to be padded. float_padding_value (float, optional): The padding value for `float` dtype tensors. non_float_padding_value (int, optional): The padding value for non-`float` (e.g., `int`) dtype tensors. max_size (int, optional): Maximum size to pad to. If None, pads to batch maximum. Returns: A tuple containing the padded value and an optional mask. """ mask = None cat_dim = None elem = values[0] dtype = elem.dtype padding_value = ( float_padding_value if torch.is_floating_point(elem) else non_float_padding_value ) if elem.dim() == 0: values = [value.unsqueeze(0) for value in values] else: if dtype != torch.float: if dtype in [torch.uint8, torch.bool]: # NOTE: We cannot use `torch.nn.utils.rnn.pad_sequence` directly with unsigned integer/boolean tensors. values = [value.long() for value in values] if "edge_index" in key and elem.dim() == 2 and elem.shape[0] == 2: # Concatenate `edge_index` tensors. # NOTE: Here, we assume dimension `0` denotes the source and target node indices and dimension `1` denotes the number of edges. if max_size is not None: # Pad each tensor to max_size before permuting padded_values = [] for val in values: val_transposed = val.permute(1, 0) # [num_edges, 2] if val_transposed.size(0) < max_size: pad_size = max_size - val_transposed.size(0) padding = torch.full( (pad_size, 2), padding_value, dtype=val_transposed.dtype, device=val_transposed.device, ) val_padded = torch.cat([val_transposed, padding], dim=0) else: val_padded = val_transposed[:max_size] padded_values.append( val_padded.unsqueeze(0) ) # [1, max_size, 2] values = [ value.permute(0, 2, 1) for value in padded_values ] # [1, 2, max_size] else: values = [ value.permute(1, 0).unsqueeze(0) for value in torch.nn.utils.rnn.pad_sequence( [val.permute(1, 0) for val in values], batch_first=True, padding_value=padding_value, ) ] else: if max_size is not None: # Custom padding to fixed max_size padded_values = [] for val in values: if val.size(0) < max_size: pad_size = max_size - val.size(0) pad_shape = (pad_size,) + val.shape[1:] padding = torch.full( pad_shape, padding_value, dtype=val.dtype, device=val.device, ) val_padded = torch.cat([val, padding], dim=0) else: val_padded = val[:max_size] padded_values.append(val_padded.unsqueeze(0)) values = padded_values else: values = [ value.unsqueeze(0) for value in torch.nn.utils.rnn.pad_sequence( values, batch_first=True, padding_value=padding_value ) ] if dtype in [torch.uint8, torch.bool]: # NOTE: We cannot use `torch.nn.utils.rnn.pad_sequence` directly with unsigned integer/boolean tensors. mask = torch.cat( [(value != padding_value) for value in values], dim=cat_dim or 0 ) for value in values: value[value == padding_value] = 0 values = [value.to(dtype) for value in values] else: if max_size is not None: # Custom padding to fixed max_size for float tensors padded_values = [] for val in values: if val.size(0) < max_size: pad_size = max_size - val.size(0) pad_shape = (pad_size,) + val.shape[1:] padding = torch.full( pad_shape, padding_value, dtype=val.dtype, device=val.device ) val_padded = torch.cat([val, padding], dim=0) else: val_padded = val[:max_size] padded_values.append(val_padded.unsqueeze(0)) values = padded_values else: values = [ value.unsqueeze(0) for value in torch.nn.utils.rnn.pad_sequence( values, batch_first=True, padding_value=padding_value ) ] return values, mask def _dense_padded_collate( key: str, values: List[Any], data_list: List[BaseData], stores: List[BaseStorage], float_padding_value: float = FLOAT_PADDING_VALUE, non_float_padding_value: int = NON_FLOAT_PADDING_VALUE, max_size: Optional[int] = None, ) -> Tuple[Any, Any]: """Collates a list of values into a single tensor or list of tensors. Args: key (str): The key of the attribute to be collated. values (list): A list of values to be collated. data_list (list): A list of data objects. stores (list): A list of storage objects. float_padding_value (float, optional): The padding value for `float` dtype tensors. non_float_padding_value (int, optional): The padding value for non-`float` (e.g., `int`) dtype tensors. max_size (int, optional): Maximum size to pad to. If None, pads to batch maximum. Returns: A tuple containing the collated value and an optional mask. """ cat_dim = None elem = values[0] if isinstance(elem, torch.Tensor) and not is_sparse(elem): # Concatenate a list of `torch.Tensor` along a new `batch_dim=0`. padding_value = ( float_padding_value if torch.is_floating_point(elem) else non_float_padding_value ) values, mask = _dense_pad_tensor( key, values, float_padding_value=float_padding_value, non_float_padding_value=non_float_padding_value, max_size=max_size, ) if getattr(elem, "is_nested", False): raise NotImplementedError( "Dense padding collation for nested tensors is not supported (tested) yet." ) tensors = [] for nested_tensor in values: tensors.extend(nested_tensor.unbind()) value = torch.nested.nested_tensor(tensors) mask = torch.nested.map(lambda tensor: (tensor != padding_value), value) return value, mask out = None if torch.utils.data.get_worker_info() is not None: # Write directly into shared memory to avoid an extra copy: numel = sum(value.numel() for value in values) if torch_geometric.typing.WITH_PT20: storage = elem.untyped_storage()._new_shared( numel * elem.element_size(), device=elem.device ) elif torch_geometric.typing.WITH_PT112: storage = elem.storage()._new_shared(numel, device=elem.device) else: storage = elem.storage()._new_shared(numel) shape = [len(data_list)] + list( values[np.argmax([value.numel() for value in values])].shape[1:] ) out = elem.new(storage).resize_(shape) value = torch.cat(values, dim=cat_dim or 0, out=out) mask = mask if mask is not None else (value != padding_value) return value, mask elif isinstance(elem, TensorFrame): raise NotImplementedError( "Dense padding collation for TensorFrames is not supported (tested) yet." ) values, mask = _dense_pad_tensor( key, values, non_float_padding_value=non_float_padding_value, max_size=max_size, ) value = torch_frame.cat(values, along="row") return value, mask elif is_sparse(elem): # Concatenate a list of `SparseTensor` along the `cat_dim`. raise NotImplementedError( "Dense padding collation for SparseTensors is not supported (tested) yet." ) values, mask = _dense_pad_tensor( key, values, non_float_padding_value=non_float_padding_value, max_size=max_size, ) if is_torch_sparse_tensor(elem): value = cat(values, dim=cat_dim) else: value = torch_sparse.cat(values, dim=cat_dim) return value, mask elif isinstance(elem, (int, float)): # Convert a list of numerical values to a `torch.Tensor`. value = torch.tensor(values) return value, None elif isinstance(elem, Mapping): # Recursively collate elements of dictionaries. value_dict, mask_dict = {}, {} for key in elem.keys(): value_dict[key], mask_dict[key] = _dense_padded_collate( key, [v[key] for v in values], data_list, stores, float_padding_value=float_padding_value, non_float_padding_value=non_float_padding_value, max_size=max_size, ) return value_dict, mask_dict elif ( isinstance(elem, Sequence) and not isinstance(elem, str) and len(elem) > 0 and isinstance(elem[0], (torch.Tensor, SparseTensor)) ): # Recursively collate elements of lists. value_list, mask_list = [], [] for i in range(len(elem)): value, mask = _dense_padded_collate( key, [v[i] for v in values], data_list, stores, float_padding_value=float_padding_value, non_float_padding_value=non_float_padding_value, max_size=max_size, ) value_list.append(value) mask_list.append(mask) return value_list, mask_list else: # Other-wise, just return the list of values as it is. return values, None def dense_padded_collate( cls, data_list: List[BaseData], follow_batch: Optional[List[str]] = None, exclude_keys: Optional[List[str]] = None, max_size: Optional[int] = None, ) -> Tuple[BaseData, Mapping]: """Collates a list of `data` objects into a single dense object of type `cls`. `collate` can handle both homogeneous and heterogeneous data objects by individually collating all their stores. In addition, `collate` can handle nested data structures such as dictionaries and lists. Args: cls (type): The class into which to collate. data_list (list): A list of data objects. follow_batch (list, optional): Creates assignment batch vectors for each key in the list. (default: `None`) exclude_keys (list, optional): Will exclude each key in the list. (default: `None`) max_size (int, optional): Maximum size to pad to. If None, pads to batch maximum. (default: `None`) Returns: A tuple containing the collated data object and a dictionary holding optional batch feature masks. """ if not isinstance(data_list, (list, tuple)): # Materialize `data_list` to keep the `_parent` weakref alive. data_list = list(data_list) if cls != data_list[0].__class__: out = cls(_base_cls=data_list[0].__class__) # Dynamic inheritance. else: out = cls() # Create empty stores: out.stores_as(data_list[0]) follow_batch = set(follow_batch or []) exclude_keys = set(exclude_keys or []) # Group all storage objects of every data object in the `data_list` by key, # i.e. `key_to_store_list = { key: [store_1, store_2, ...], ... }`: key_to_stores = defaultdict(list) for data in data_list: for store in data.stores: key_to_stores[store._key].append(store) # With this, we iterate over each list of storage objects and recursively # collate all its attributes into a unified representation: # We maintain an additional dictionary: # * `mask_dict` stores a mask representation of each attribute # and is needed to re-construct individual elements from mini-batches. mask_dict = defaultdict(dict) for out_store in out.stores: key = out_store._key stores = key_to_stores[key] for attr in stores[0].keys(): if attr in exclude_keys: # Do not include top-level attribute. continue values = [store[attr] for store in stores] # The `num_nodes` attribute needs special treatment, as we need to # sum their values up instead of merging them to a list: if attr == "num_nodes": out_store._num_nodes = values out_store.num_nodes = sum(values) continue # Skip batching of `ptr` vectors during dense collation: if attr == "ptr": continue # Collate attributes into a unified representation: value, mask = _dense_padded_collate( attr, values, data_list, stores, max_size=max_size ) out_store[attr] = value if mask is not None: if key is not None: mask_dict[key][attr] = mask else: mask_dict[attr] = mask return out, mask_dict def dense_padded_from_data_list( data_list: List[BaseData], follow_batch: Optional[List[str]] = None, exclude_keys: Optional[List[str]] = None, max_size: Optional[int] = None, ): r"""Constructs a dense `~torch_geometric.data.Batch` object from a Python list of `~torch_geometric.data.Data` or `~torch_geometric.data.HeteroData` objects. The assignment vector `batch` is created on the fly. In addition, creates assignment vectors for each key in `follow_batch`. Will exclude any keys given in `exclude_keys`. Args: data_list (list): A list of data objects. follow_batch (list, optional): Creates assignment batch vectors for each key in the list. (default: `None`) exclude_keys (list, optional): Will exclude each key in the list. (default: `None`) max_size (int, optional): Maximum size to pad to. If None, pads to batch maximum. (default: `None`) Returns: A single `Batch` object holding a mini-batch of data. """ batch, mask_dict = dense_padded_collate( Batch, data_list=data_list, follow_batch=follow_batch, exclude_keys=exclude_keys, max_size=max_size, ) batch._num_graphs = len(data_list) batch.mask_dict = mask_dict # NOTE: `mask_dict` must be a public attribute of `Batch` for auto-device onloading to work. return batch class DensePaddingCollater: """Collates data objects from a `torch_geometric.data.Dataset` to a mini-batch using padding along with a new dimension. Data objects can be either of type `~torch_geometric.data.Data` or `~torch_geometric.data.HeteroData`. """ def __init__( self, dataset: Union[Dataset, Sequence[BaseData], DatasetAdapter], follow_batch: Optional[List[str]] = None, exclude_keys: Optional[List[str]] = None, max_size: Optional[int] = None, ): self.dataset = dataset self.follow_batch = follow_batch self.exclude_keys = exclude_keys self.max_size = max_size def __call__(self, batch: List[Any]) -> Any: """Collates a python list of data objects to the internal storage format of `torch_geometric.data.DataLoader`. Args: batch: A python list of data objects. Returns: A mini-batch of data objects. """ elem = batch[0] if isinstance(elem, BaseData): return dense_padded_from_data_list( batch, follow_batch=self.follow_batch, exclude_keys=self.exclude_keys, max_size=self.max_size, ) elif isinstance(elem, torch.Tensor): return default_collate(batch) elif isinstance(elem, TensorFrame): return torch_frame.cat(batch, along="row") elif isinstance(elem, float): return torch.tensor(batch, dtype=torch.float) elif isinstance(elem, int): return torch.tensor(batch) elif isinstance(elem, str): return batch elif isinstance(elem, Mapping): return {key: self([data[key] for data in batch]) for key in elem} elif isinstance(elem, tuple) and hasattr(elem, "_fields"): return type(elem)(*(self(s) for s in zip(*batch))) elif isinstance(elem, Sequence) and not isinstance(elem, str): return [self(s) for s in zip(*batch)] raise TypeError(f"DataLoader found invalid type: '{type(elem)}'") def collate_fn(self, batch: List[Any]) -> Any: """Collates a python list of data objects to the internal storage format of `torch_geometric.data.DataLoader`. Args: batch: A python list of data objects. Returns: A mini-batch of data objects. """ if isinstance(self.dataset, OnDiskDataset): return self(self.dataset.multi_get(batch)) return self(batch) class DensePaddingDataLoader(torch.utils.data.DataLoader): r"""A data loader which merges data objects from a `torch_geometric.data.Dataset` to a mini- batch using padding along with a new dimension. Data objects can be either of type `~torch_geometric.data.Data` or `~torch_geometric.data.HeteroData`. Args: dataset (Dataset): The dataset from which to load the data. batch_size (int, optional): How many samples per batch to load. (default: `1`) shuffle (bool, optional): If set to `True`, the data will be reshuffled at every epoch. (default: `False`) follow_batch (List[str], optional): Creates assignment batch vectors for each key in the list. (default: `None`) exclude_keys (List[str], optional): Will exclude each key in the list. (default: `None`) max_size (int, optional): Maximum size to pad to. If None, pads to batch maximum. (default: `None`) **kwargs (optional): Additional arguments of `torch.utils.data.DataLoader`. """ def __init__( self, dataset: Union[Dataset, Sequence[BaseData], DatasetAdapter], batch_size: int = 1, shuffle: bool = False, follow_batch: Optional[List[str]] = None, exclude_keys: Optional[List[str]] = None, max_size: Optional[int] = None, **kwargs, ): # Remove for PyTorch Lightning: kwargs.pop("collate_fn", None) # Save for PyTorch Lightning < 1.6: self.follow_batch = follow_batch self.exclude_keys = exclude_keys self.max_size = max_size self.collator = DensePaddingCollater( dataset, follow_batch, exclude_keys, max_size ) if isinstance(dataset, OnDiskDataset): dataset = range(len(dataset)) super().__init__( dataset, batch_size, shuffle, collate_fn=self.collator.collate_fn, **kwargs, )