ContextQMD
Back to Claude Scientific Skills

Datasets

skills/torchdrug/references/datasets.md

2.54.05.3 KB
Original Source

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.

python
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
python
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
python
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:

python
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:

python
train_set, valid_set, test_set = dataset.split()

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

python
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():

python
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:

python
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.