ProteinNet documents 40 test proteins, but the pinned archive contains 34

[adpartin]

Issue

torchdrug/datasets/proteinnet.py documents the following split sizes in the ProteinNet class docstring:

Statistics:
    - #Train: 25,299
    - #Valid: 224
    - #Test: 40

The same class pins the archive it downloads:

url = "https://miladeepgraphlearningproteindata.s3.us-east-2.amazonaws.com/data/proteinnet.tar.gz"
md5 = "ab44ab201b1570c0171a2bba9eb4d389"

That archive contains 34 test proteins, not 40. The train and valid counts are correct.

Reproduction

import os
import pickle

import lmdb
from torchdrug import datasets, utils

path = os.path.expanduser("~/torchdrug-datasets/")
os.makedirs(path, exist_ok=True)

cls = datasets.ProteinNet
tar = utils.download(cls.url, path, md5=cls.md5)
data_dir = utils.extract(tar)

print("archive md5:", utils.compute_md5(tar))
for split in cls.splits:
    file_path = os.path.join(
        data_dir, "proteinnet/proteinnet_%s.lmdb" % split
    )
    with lmdb.open(file_path, readonly=True, lock=False).begin() as txn:
        print(split, pickle.loads(txn.get(b"num_examples")))

Output:

archive md5: ab44ab201b1570c0171a2bba9eb4d389
train 25299
valid 224
test 34

The ProteinNet source file included in the PyPI 0.2.1 release is identical to the file currently on master (checked 2026-09-06).

Comparison with TAPE

TAPE's download_data.sh downloads ProteinNet from:

http://s3.amazonaws.com/songlabdata/proteindata/data_pytorch/proteinnet.tar.gz

Fetching that file gives MD5:

85c055d47487df2725e422599d9465d9

TorchDrug vs TAPE comparison of the proteinnet downloads:

split	TorchDrug archive	TAPE archive	ID comparison
train	25,299	25,299	identical
valid	224	224	identical
test	34	40	0 shared IDs

The train and valid IDs match. The two test splits share no IDs, so this is not the same test set with 6 records removed.

Every TAPE test ID appears in the official CASP12 target list. Every TorchDrug test ID appears in the official CASP14 target list. The documented count of 40 is therefore consistent with the CASP12 test set, while the pinned archive contains the 34-protein CASP14 test set.

Related

The PEER benchmark contact-prediction task reports 40 test proteins and states that it uses the CASP12 test set. Its contact-prediction configs load TorchDrug's ProteinNet class. A run with a clean dataset cache therefore downloads the 34-protein CASP14 test set.

Suggested fix

If the pinned archive is intended, please update the documented test count in the ProteinNet docstring and name the CASP version. If the 40-protein CASP12 split is intended, please update the URL and MD5 to point to an archive containing that split.