torch_geometric.data.Dataset

class Dataset(root: Optional[str] = None, transform: Optional[Callable] = None, pre_transform: Optional[Callable] = None, pre_filter: Optional[Callable] = None, log: bool = True, force_reload: bool = False)[source]

Bases: Dataset

Dataset base class for creating graph datasets. See here for the accompanying tutorial.

Parameters:

  • root (str, optional) – Root directory where the dataset should be saved. (optional: None)

  • transform (callable, optional) – A function/transform that takes in a Data or HeteroData object and returns a transformed version. The data object will be transformed before every access. (default: None)

  • pre_transform (callable, optional) – A function/transform that takes in a Data or HeteroData object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)

  • pre_filter (callable, optional) – A function that takes in a Data or HeteroData object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)

  • log (bool, optional) – Whether to print any console output while downloading and processing the dataset. (default: True)

  • force_reload (bool, optional) – Whether to re-process the dataset. (default: False)

property raw_file_names: Union[str, List[str], Tuple[str, ...]]

The name of the files in the self.raw_dir folder that must be present in order to skip downloading.

property processed_file_names: Union[str, List[str], Tuple[str, ...]]

The name of the files in the self.processed_dir folder that must be present in order to skip processing.

download() → None[source]

Downloads the dataset to the self.raw_dir folder.

process() → None[source]

Processes the dataset to the self.processed_dir folder.

len() → int[source]

Returns the number of data objects stored in the dataset.

get(idx: int) → BaseData[source]

Gets the data object at index idx.

property num_node_features: int

Returns the number of features per node in the dataset.

property num_features: int

Returns the number of features per node in the dataset. Alias for num_node_features.

property num_edge_features: int

Returns the number of features per edge in the dataset.

property num_classes: int

Returns the number of classes in the dataset.

property raw_paths: List[str]

The absolute filepaths that must be present in order to skip downloading.

property processed_paths: List[str]

The absolute filepaths that must be present in order to skip processing.

property has_download: bool

Checks whether the dataset defines a download() method.

property has_process: bool

Checks whether the dataset defines a process() method.

index_select(idx: Union[slice, Tensor, ndarray, Sequence]) → Dataset[source]

Creates a subset of the dataset from specified indices idx. Indices idx can be a slicing object, e.g., [2:5], a list, a tuple, or a torch.Tensor or np.ndarray of type long or bool.

shuffle(return_perm: bool = False) → Union[Dataset, Tuple[Dataset, Tensor]][source]

Randomly shuffles the examples in the dataset.

Parameters:

return_perm (bool, optional) – If set to True, will also return the random permutation used to shuffle the dataset. (default: False)

get_summary() → Any[source]

Collects summary statistics for the dataset.

print_summary() → None[source]

Prints summary statistics of the dataset to the console.

to_datapipe() → Any[source]

Converts the dataset into a torch.utils.data.DataPipe.

The returned instance can then be used with built-in DataPipes for baching graphs as follows:

from torch_geometric.datasets import QM9

dp = QM9(root='./data/QM9/').to_datapipe()
dp = dp.batch_graphs(batch_size=2, drop_last=True)

for batch in dp:
    pass

See the PyTorch tutorial for further background on DataPipes.