torch_geometric.datasets

class KarateClub(transform=None)[source]

Zachary’s karate club network from the “An Information Flow Model for Conflict and Fission in Small Groups” paper, containing 34 nodes, connected by 154 (undirected and unweighted) edges. Every node is labeled by one of two classes.

Parameters:transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data object and returns a transformed version. The data object will be transformed before every access. (default: None)
class TUDataset(root, name, transform=None, pre_transform=None, pre_filter=None, use_node_attr=False, use_edge_attr=False, cleaned=False)[source]

A variety of graph kernel benchmark datasets, .e.g. “IMDB-BINARY”, “REDDIT-BINARY” or “PROTEINS”, collected from the TU Dortmund University. In addition, this dataset wrapper provides cleaned dataset versions as motivated by the “Understanding Isomorphism Bias in Graph Data Sets” paper, containing only non-isomorphic graphs.

Note

Some datasets may not come with any node labels. You can then either make use of the argument use_node_attr to load additional continuous node attributes (if present) or provide synthetic node features using transforms such as like torch_geometric.transforms.Constant or torch_geometric.transforms.OneHotDegree.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • name (string) – The name of the dataset.
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
  • use_node_attr (bool, optional) – If True, the dataset will contain additional continuous node attributes (if present). (default: False)
  • use_edge_attr (bool, optional) – If True, the dataset will contain additional continuous edge attributes (if present). (default: False)
  • cleaned – (bool, optional): If True, the dataset will contain only non-isomorphic graphs. (default: False)
class Planetoid(root, name, transform=None, pre_transform=None)[source]

The citation network datasets “Cora”, “CiteSeer” and “PubMed” from the “Revisiting Semi-Supervised Learning with Graph Embeddings” paper. Nodes represent documents and edges represent citation links. Training, validation and test splits are given by binary masks.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • name (string) – The name of the dataset ("Cora", "CiteSeer", "PubMed").
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)
class CitationFull(root, name, transform=None, pre_transform=None)[source]

The full citation network datasets from the “Deep Gaussian Embedding of Graphs: Unsupervised Inductive Learning via Ranking” paper. Nodes represent documents and edges represent citation links. Datasets include citeseer, cora, cora_ml, dblp, pubmed.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • name (string) – The name of the dataset ("Cora", "Cora_ML" "CiteSeer", "DBLP", "PubMed").
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)
class CoraFull(root, transform=None, pre_transform=None)[source]

Alias for torch_geometric.dataset.CitationFull with name="cora".

class Coauthor(root, name, transform=None, pre_transform=None)[source]

The Coauthor CS and Coauthor Physics networks from the “Pitfalls of Graph Neural Network Evaluation” paper. Nodes represent authors that are connected by an edge if they co-authored a paper. Given paper keywords for each author’s papers, the task is to map authors to their respective field of study.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • name (string) – The name of the dataset ("CS", "Physics").
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)
class Amazon(root, name, transform=None, pre_transform=None)[source]

The Amazon Computers and Amazon Photo networks from the “Pitfalls of Graph Neural Network Evaluation” paper. Nodes represent goods and edges represent that two goods are frequently bought together. Given product reviews as bag-of-words node features, the task is to map goods to their respective product category.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • name (string) – The name of the dataset ("Computers", "Photo").
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)
class PPI(root, split='train', transform=None, pre_transform=None, pre_filter=None)[source]

The protein-protein interaction networks from the “Predicting Multicellular Function through Multi-layer Tissue Networks” paper, containing positional gene sets, motif gene sets and immunological signatures as features (50 in total) and gene ontology sets as labels (121 in total).

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • split (string) – If "train", loads the training dataset. If "val", loads the validation dataset. If "test", loads the test dataset. (default: "train")
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class Reddit(root, transform=None, pre_transform=None, pre_filter=None)[source]

The Reddit dataset from the “Inductive Representation Learning on Large Graphs” paper, containing Reddit posts belonging to different communities.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class QM7b(root, transform=None, pre_transform=None, pre_filter=None)[source]

The QM7b dataset from the “MoleculeNet: A Benchmark for Molecular Machine Learning” paper, consisting of 7,211 molecules with 14 regression targets.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class QM9(root, transform=None, pre_transform=None, pre_filter=None)[source]

The QM9 dataset from the “MoleculeNet: A Benchmark for Molecular Machine Learning” paper, consisting of about 130,000 molecules with 16 regression targets. Each molecule includes complete spatial information for the single low energy conformation of the atoms in the molecule. In addition, we provide the atom features from the “Neural Message Passing for Quantum Chemistry” paper.

Target Property Description Unit
0 \(\mu\) Dipole moment \(\textrm{D}\)
1 \(\alpha\) Isotropic polarizability \({a_0}^3\)
2 \(\epsilon_{\textrm{HOMO}}\) Highest occupied molecular orbital energy \(E_{\textrm{h}}\)
3 \(\epsilon_{\textrm{LUMO}}\) Lowest unoccupied molecular orbital energy \(E_{\textrm{h}}\)
4 \(\Delta \epsilon\) Gap between \(\epsilon_{\textrm{HOMO}}\) and \(\epsilon_{\textrm{LUMO}}\) \(E_{\textrm{h}}\)
5 \(\langle R^2 \rangle\) Electronic spatial extent \({a_0}^2\)
6 \(\textrm{ZPVE}\) Zero point vibrational energy \(E_{\textrm{h}}\)
7 \(U_0\) Internal energy at 0K \(E_{\textrm{h}}\)
8 \(U\) Internal energy at 298.15K \(E_{\textrm{h}}\)
9 \(H\) Enthalpy at 298.15K \(E_{\textrm{h}}\)
10 \(G\) Free energy at 298.15K \(E_{\textrm{h}}\)
11 \(c_{\textrm{v}}\) Heat capavity at 298.15K \(\frac{\textrm{cal}}{\textrm{mol K}}\)
12 \(U_0^{\textrm{ATOM}}\) Atomization energy at 0K \(\frac{\textrm{kcal}}{\textrm{mol}}\)
13 \(U^{\textrm{ATOM}}\) Atomization energy at 298.15K \(\frac{\textrm{kcal}}{\textrm{mol}}\)
14 \(H^{\textrm{ATOM}}\) Atomization enthalpy at 298.15K \(\frac{\textrm{kcal}}{\textrm{mol}}\)
15 \(G^{\textrm{ATOM}}\) Atomization free energy at 298.15K \(\frac{\textrm{kcal}}{\textrm{mol}}\)
Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class Entities(root, name, transform=None, pre_transform=None)[source]

The relational entities networks “AIFB”, “MUTAG”, “BGS” and “AM” from the “Modeling Relational Data with Graph Convolutional Networks” paper. Training and test splits are given by node indices.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • name (string) – The name of the dataset ("AIFB", "MUTAG", "BGS", "AM").
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)
class GEDDataset(root, name, train=True, transform=None, pre_transform=None, pre_filter=None)[source]

The GED datasets from the “Graph Edit Distance Computation via Graph Neural Networks” paper. GEDs can be accessed via the global attributes ged and norm_ged and provide the exact GEDs for all train/train graph pairs and all train/test graph pairs.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • name (string) – The name of the dataset (one of "AIDS700nef", "LINUX", "ALKANE", "IMDBMulti").
  • train (bool, optional) – If True, loads the training dataset, otherwise the test dataset. (default: True)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class MNISTSuperpixels(root, train=True, transform=None, pre_transform=None, pre_filter=None)[source]

MNIST superpixels dataset from the “Geometric Deep Learning on Graphs and Manifolds Using Mixture Model CNNs” paper, containing 70,000 graphs with 75 nodes each. Every graph is labeled by one of 10 classes.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • train (bool, optional) – If True, loads the training dataset, otherwise the test dataset. (default: True)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class FAUST(root, train=True, transform=None, pre_transform=None, pre_filter=None)[source]

The FAUST humans dataset from the “FAUST: Dataset and Evaluation for 3D Mesh Registration” paper, containing 100 watertight meshes representing 10 different poses for 10 different subjects.

Note

Data objects hold mesh faces instead of edge indices. To convert the mesh to a graph, use the torch_geometric.transforms.FaceToEdge as pre_transform. To convert the mesh to a point cloud, use the torch_geometric.transforms.SamplePoints as transform to sample a fixed number of points on the mesh faces according to their face area.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • train (bool, optional) – If True, loads the training dataset, otherwise the test dataset. (default: True)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class DynamicFAUST(root, subjects=None, categories=None, transform=None, pre_transform=None, pre_filter=None)[source]

The dynamic FAUST humans dataset from the “Dynamic FAUST: Registering Human Bodies in Motion” paper.

Note

Data objects hold mesh faces instead of edge indices. To convert the mesh to a graph, use the torch_geometric.transforms.FaceToEdge as pre_transform. To convert the mesh to a point cloud, use the torch_geometric.transforms.SamplePoints as transform to sample a fixed number of points on the mesh faces according to their face area.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • subjects (list, optional) – List of subjects to include in the dataset. Can include the subjects "50002", "50004", "50007", "50009", "50020", "50021", "50022", "50025", "50026", "50027". If set to None, the dataset will contain all subjects. (default: None)
  • categories (list, optional) – List of categories to include in the dataset. Can include the categories "chicken_wings", "hips", "jiggle_on_toes", "jumping_jacks", "knees", "light_hopping_loose", "light_hopping_stiff", "one_leg_jump", "one_leg_loose", "personal_move", "punching", "running_on_spot", "running_on_spot_bugfix", "shake_arms", "shake_hips", "shoulders". If set to None, the dataset will contain all categories. (default: None)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class ShapeNet(root, categories=None, include_normals=True, split='trainval', transform=None, pre_transform=None, pre_filter=None)[source]

The ShapeNet part level segmentation dataset from the “A Scalable Active Framework for Region Annotation in 3D Shape Collections” paper, containing about 17,000 3D shape point clouds from 16 shape categories. Each category is annotated with 2 to 6 parts.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • categories (string or [string], optional) – The category of the CAD models (one or a combination of "Airplane", "Bag", "Cap", "Car", "Chair", "Earphone", "Guitar", "Knife", "Lamp", "Laptop", "Motorbike", "Mug", "Pistol", "Rocket", "Skateboard", "Table"). Can be explicitly set to None to load all categories. (default: None)
  • include_normals (bool, optional) – If set to False, will not include normal vectors as input features. (default: True)
  • split (string, optional) – If "train", loads the training dataset. If "val", loads the validation dataset. If "trainval", loads the training and validation dataset. If "test", loads the test dataset. (default: "trainval")
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class ModelNet(root, name='10', train=True, transform=None, pre_transform=None, pre_filter=None)[source]

The ModelNet10/40 datasets from the “3D ShapeNets: A Deep Representation for Volumetric Shapes” paper, containing CAD models of 10 and 40 categories, respectively.

Note

Data objects hold mesh faces instead of edge indices. To convert the mesh to a graph, use the torch_geometric.transforms.FaceToEdge as pre_transform. To convert the mesh to a point cloud, use the torch_geometric.transforms.SamplePoints as transform to sample a fixed number of points on the mesh faces according to their face area.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • name (string, optional) – The name of the dataset ("10" for ModelNet10, "40" for ModelNet40). (default: "10")
  • train (bool, optional) – If True, loads the training dataset, otherwise the test dataset. (default: True)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class CoMA(root, train=True, transform=None, pre_transform=None, pre_filter=None)[source]

The CoMA 3D faces dataset from the “Generating 3D faces using Convolutional Mesh Autoencoders” paper, containing 20,466 meshes of extreme expressions captured over 12 different subjects.

Note

Data objects hold mesh faces instead of edge indices. To convert the mesh to a graph, use the torch_geometric.transforms.FaceToEdge as pre_transform. To convert the mesh to a point cloud, use the torch_geometric.transforms.SamplePoints as transform to sample a fixed number of points on the mesh faces according to their face area.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • train (bool, optional) – If True, loads the training dataset, otherwise the test dataset. (default: True)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class SHREC2016(root, partiality, category, train=True, transform=None, pre_transform=None, pre_filter=None)[source]

The SHREC 2016 partial matching dataset from the “SHREC‘16: Partial Matching of Deformable Shapes” paper. The reference shape can be referenced via dataset.ref.

Note

Data objects hold mesh faces instead of edge indices. To convert the mesh to a graph, use the torch_geometric.transforms.FaceToEdge as pre_transform. To convert the mesh to a point cloud, use the torch_geometric.transforms.SamplePoints as transform to sample a fixed number of points on the mesh faces according to their face area.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • partiality (string) – The partiality of the dataset (one of "Holes", "Cuts").
  • category (string) – The category of the dataset (one of "Cat", "Centaur", "David", "Dog", "Horse", "Michael", "Victoria", "Wolf").
  • train (bool, optional) – If True, loads the training dataset, otherwise the test dataset. (default: True)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class TOSCA(root, categories=None, transform=None, pre_transform=None, pre_filter=None)[source]

The TOSCA dataset from the “Numerical Geometry of Non-Ridig Shapes” book, containing 80 meshes. Meshes within the same category have the same triangulation and an equal number of vertices numbered in a compatible way.

Note

Data objects hold mesh faces instead of edge indices. To convert the mesh to a graph, use the torch_geometric.transforms.FaceToEdge as pre_transform. To convert the mesh to a point cloud, use the torch_geometric.transforms.SamplePoints as transform to sample a fixed number of points on the mesh faces according to their face area.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • categories (list, optional) – List of categories to include in the dataset. Can include the categories "Cat", "Centaur", "David", "Dog", "Gorilla", "Horse", "Michael", "Victoria", "Wolf". If set to None, the dataset will contain all categories. (default: None)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class PCPNetDataset(root, category, split='train', transform=None, pre_transform=None, pre_filter=None)[source]

The PCPNet dataset from the “PCPNet: Learning Local Shape Properties from Raw Point Clouds” paper, consisting of 30 shapes, each given as a point cloud, densely sampled with 100k points. For each shape, surface normals and local curvatures are given as node features.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • category (string) – The training set category (one of "NoNoise", "Noisy", "VarDensity", "NoisyAndVarDensity" for split="train" or split="val", or one of "All", "LowNoise", "MedNoise", "HighNoise", :obj:”VarDensityStriped”, "VarDensityGradient" for split="test").
  • split (string) – If "train", loads the training dataset. If "val", loads the validation dataset. If "test", loads the test dataset. (default: "train")
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class S3DIS(root, test_area=6, train=True, transform=None, pre_transform=None, pre_filter=None)[source]

The (pre-processed) Stanford Large-Scale 3D Indoor Spaces dataset from the “3D Semantic Parsing of Large-Scale Indoor Spaces” paper, containing point clouds of six large-scale indoor parts in three buildings with 12 semantic elements (and one clutter class).

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • test_area (int, optional) – Which area to use for testing (1-6). (default: 6)
  • train (bool, optional) – If True, loads the training dataset, otherwise the test dataset. (default: True)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class GeometricShapes(root, train=True, transform=None, pre_transform=None, pre_filter=None)[source]

Synthetic dataset of various geometric shapes like cubes, spheres or pyramids.

Note

Data objects hold mesh faces instead of edge indices. To convert the mesh to a graph, use the torch_geometric.transforms.FaceToEdge as pre_transform. To convert the mesh to a point cloud, use the torch_geometric.transforms.SamplePoints as transform to sample a fixed number of points on the mesh faces according to their face area.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • train (bool, optional) – If True, loads the training dataset, otherwise the test dataset. (default: True)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class BitcoinOTC(root, edge_window_size=10, transform=None, pre_transform=None)[source]

The Bitcoin-OTC dataset from the “EvolveGCN: Evolving Graph Convolutional Networks for Dynamic Graphs” paper, consisting of 138 who-trusts-whom networks of sequential time steps.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • edge_window_size (int, optional) – The window size for the existence of an edge in the graph sequence since its initial creation. (default: 10)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)
class ICEWS18(root, split='train', transform=None, pre_transform=None, pre_filter=None)[source]

The Integrated Crisis Early Warning System (ICEWS) dataset used in the, e.g., “Recurrent Event Network for Reasoning over Temporal Knowledge Graphs” paper, consisting of events collected from 1/1/2018 to 10/31/2018 (24 hours time granularity).

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • split (string) – If "train", loads the training dataset. If "val", loads the validation dataset. If "test", loads the test dataset. (default: "train")
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class GDELT(root, split='train', transform=None, pre_transform=None, pre_filter=None)[source]

The Global Database of Events, Language, and Tone (GDELT) dataset used in the, e.g., “Recurrent Event Network for Reasoning over Temporal Knowledge Graphs” paper, consisting of events collected from 1/1/2018 to 1/31/2018 (15 minutes time granularity).

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • split (string) – If "train", loads the training dataset. If "val", loads the validation dataset. If "test", loads the test dataset. (default: "train")
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class DBP15K(root, pair, transform=None, pre_transform=None)[source]

The DBP15K dataset from the “Cross-lingual Entity Alignment via Joint Attribute-Preserving Embedding” paper, where Chinese, Japanese and French versions of DBpedia were linked to its English version. Node features are given by pre-trained and aligned monolingual word embeddings from the “Cross-lingual Knowledge Graph Alignment via Graph Matching Neural Network” paper.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • pair (string) – The pair of languages ("en_zh", "en_fr", "en_ja", "zh_en", "fr_en", "ja_en").
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)
class WILLOWObjectClass(root, category, transform=None, pre_transform=None, pre_filter=None)[source]

The WILLOW-ObjectClass dataset from the “Learning Graphs to Match” paper, containing 10 equal keypoints of at least 40 images in each category. The keypoints contain interpolated features from a pre-trained VGG16 model on ImageNet (relu4_2 and relu5_1).

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • category (string) – The category of the images (one of "Car", "Duck", "Face", "Motorbike", "Winebottle")
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class PascalVOCKeypoints(root, category, train=True, transform=None, pre_transform=None, pre_filter=None)[source]

The Pascal VOC 2011 dataset with Berkely annotations of keypoints from the “Poselets: Body Part Detectors Trained Using 3D Human Pose Annotations” paper, containing 0 to 23 keypoints per example over 20 categories. The dataset is pre-filtered to exclude difficult, occluded and truncated objects. The keypoints contain interpolated features from a pre-trained VGG16 model on ImageNet (relu4_2 and relu5_1).

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • category (string) – The category of the images (one of "Aeroplane", "Bicycle", "Bird", "Boat", "Bottle", "Bus", "Car", "Cat", "Chair", "Diningtable", "Dog", "Horse", "Motorbike", "Person", "Pottedplant", "Sheep", "Sofa", "Train", "TVMonitor")
  • train (bool, optional) – If True, loads the training dataset, otherwise the test dataset. (default: True)
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class PascalPF(root, category, transform=None, pre_transform=None, pre_filter=None)[source]

The Pascal-PF dataset from the “Proposal Flow” paper, containing 4 to 16 keypoints per example over 20 categories.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • category (string) – The category of the images (one of "Aeroplane", "Bicycle", "Bird", "Boat", "Bottle", "Bus", "Car", "Cat", "Chair", "Diningtable", "Dog", "Horse", "Motorbike", "Person", "Pottedplant", "Sheep", "Sofa", "Train", "TVMonitor")
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class SNAPDataset(root, name, transform=None, pre_transform=None, pre_filter=None)[source]

A variety of graph datasets collected from SNAP at Stanford University.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • name (string) – The name of the dataset.
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a boolean value, indicating whether the data object should be included in the final dataset. (default: None)
class SuiteSparseMatrixCollection(root, group, name, transform=None, pre_transform=None)[source]

A suite of sparse matrix benchmarks known as the Suite Sparse Matrix Collection collected from a wide range of applications.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • group (string) – The group of the sparse matrix.
  • name (string) – The name of the sparse matrix.
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data 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 an torch_geometric.data.Data object and returns a transformed version. The data object will be transformed before being saved to disk. (default: None)
class TrackMLParticleTrackingDataset(root, transform=None)[source]

The TrackML Particle Tracking Challenge dataset to reconstruct particle tracks from 3D points left in the silicon detectors.

Parameters:
  • root (string) – Root directory where the dataset should be saved.
  • transform (callable, optional) – A function/transform that takes in an torch_geometric.data.Data object and returns a transformed version. The data object will be transformed before every access. (default: None)
get(idx)[source]

Gets the data object at index idx.