Datasets

Use the TorchDrug 0.2.1 dataset reference as the class inventory and signature source. Dataset constructors download and cache data under the path supplied by the caller.

Dataset families

Molecule property prediction

Documented classes include:

Classification: BACE, BBBP, ClinTox, HIV, MUV, SIDER, Tox21, ToxCast
Regression / quantum properties: FreeSolv, Lipophilicity, QM8, QM9, PCQM4M
Pretraining / generation: ChEMBLFiltered, ZINC250k, ZINC2m, MOSES

The official property tutorial uses ClinTox; the pretraining tutorial uses ClinTox for a small demonstration and recommends larger data such as ZINC2m for real pretraining; the generation tutorial uses ZINC250k.

from torchdrug import datasets

dataset = datasets.ClinTox(
    "~/molecule-datasets/",
    atom_feature="default",
    bond_feature="default",
)
print(dataset.tasks)
print(dataset.node_feature_dim)
print(dataset.edge_feature_dim)

Common molecule options include atom_feature, bond_feature, mol_feature, with_hydrogen, and kekulize. Availability varies by class; inspect the class signature before adding options.

Protein properties and structure

Documented families include:

Sequence / property: BetaLactamase, BinaryLocalization, SubcellularLocalization
Structure / function: EnzymeCommission, GeneOntology, AlphaFoldDB
Structure labels: Fold, SecondaryStructure
Protein-protein: HumanPPI, YeastPPI, PPIAffinity
Protein-ligand: BindingDB, PDBBind

dataset = datasets.EnzymeCommission(
    "~/protein-datasets/",
    atom_feature=None,
    bond_feature=None,
    residue_feature="default",
)
train_set, valid_set, test_set = dataset.split()

Protein datasets can be expensive to parse. Where supported, lazy=True trades lower startup memory for slower item loading. For sequence-only models, omitting atom and bond features avoids unnecessary atom-level construction.

Knowledge graphs

Documented classes:

FB15k
FB15k237
WN18
WN18RR
Hetionet

dataset = datasets.FB15k237("~/kg-datasets/")
train_set, valid_set, test_set = dataset.split()

print(dataset.num_entity)
print(dataset.num_relation)

These datasets provide predefined benchmark splits. Preserve those splits for comparable evaluation.

Retrosynthesis

USPTO50k contains 50,017 reactions across 10 reaction classes. The official G2Gs workflow loads two views:

reaction_dataset = datasets.USPTO50k(
    "~/molecule-datasets/",
    atom_feature="center_identification",
    kekulize=True,
)
synthon_dataset = datasets.USPTO50k(
    "~/molecule-datasets/",
    as_synthon=True,
    atom_feature="synthon_completion",
    kekulize=True,
)

Reaction mode yields reactant/product pairs for center identification. Synthon mode yields reactant/synthon pairs for synthon completion.

Splitting correctly

Some benchmark datasets expose predefined splits:

train_set, valid_set, test_set = dataset.split()

For the property-prediction tutorial's random 80/10/10 split, use PyTorch:

import torch

lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
lengths.append(len(dataset) - sum(lengths))
train_set, valid_set, test_set = torch.utils.data.random_split(dataset, lengths)

Do not assume dataset.split([0.8, 0.1, 0.1]) is a documented universal API.

For paired retrosynthesis views, reset the same seed before each split():

torch.manual_seed(1)
reaction_train, reaction_valid, reaction_test = reaction_dataset.split()
torch.manual_seed(1)
synthon_train, synthon_valid, synthon_test = synthon_dataset.split()

This preserves sample alignment.

Feature configuration

Dataset dimensions depend on feature choices. Construct models from the loaded dataset rather than hard-coding dimensions:

model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[256, 256, 256],
    edge_input_dim=dataset.edge_feature_dim,
)

Generation and retrosynthesis often require specialized feature sets:

Pretraining: atom_feature="pretrain", bond_feature="pretrain"
GCPN / GraphAF: atom_feature="symbol", kekulize=True
Center identification: atom_feature="center_identification"
Synthon completion: atom_feature="synthon_completion"

Do not mix checkpoint weights across incompatible feature configurations.

Data integrity and evaluation

Cache datasets in a controlled project or user data directory.
Record TorchDrug version, feature arguments, split method, and random seed.
Preserve predefined KG splits.
For molecular benchmarks, use the split protocol required by the benchmark; do not claim a random split is a scaffold split.
Inspect downloaded data licenses and provenance before redistribution.
Validate labels, missing-value masks, and task names before training.