def from_data_list(cls,
data_list: List[BaseData],
follow_batch: Optional[List[str]] = None,
exclude_keys: Optional[List[str]] = None):
r"""Constructs a :class:`~torch_geometric.data.Batch` object from a
Python list of :class:`~torch_geometric.data.Data` or
:class:`~torch_geometric.data.HeteroData` objects.
The assignment vector :obj:`batch` is created on the fly.
In addition, creates assignment vectors for each key in
:obj:`follow_batch`.
Will exclude any keys given in :obj:`exclude_keys`."""
batch, slice_dict, inc_dict = collate(
cls,
data_list=data_list,
increment=True,
add_batch=not isinstance(data_list[0], Batch),
follow_batch=follow_batch,
exclude_keys=exclude_keys,
)
batch._num_graphs = len(data_list)
batch._slice_dict = slice_dict
batch._inc_dict = inc_dict
return batch
def collate(data_list):
r"""Collates a Python list of :obj:`torch_geometric.data.Data` objects
to the internal storage format of
:class:`~torch_geometric.data.InMemoryDataset`."""
if len(data_list) == 1:
return data_list[0], None
data, slices, _ = collate(data_list[0].__class__,
data_list=data_list,
increment=False,
add_batch=False,
follow_batch=['x_i', 'x_j'])
return data, slices
def collate(
data_list: List[Data]) -> Tuple[Data, Optional[Dict[str, Tensor]]]:
r"""Collates a Python list of :obj:`torch_geometric.data.Data` objects
to the internal storage format of
:class:`~torch_geometric.data.InMemoryDataset`."""
if len(data_list) == 1:
return data_list[0], None
data, slices, _ = collate(
data_list[0].__class__,
data_list=data_list,
increment=False,
add_batch=False,
)
return data, slices