Reproducing this work end to end¶
Every number in reports/ comes from the commands below. Nothing here depends
on the specific cluster used during development: any machine with a CUDA GPU
works for training, and the evaluation half runs on a laptop.
0. Environment¶
Two environments are needed, because MolScribe (the external baseline) pins
torch<2 and cannot share an interpreter with a current PyTorch.
Training / evaluation (Python 3.11+):
python -m venv .venv && . .venv/bin/activate
pip install -e ".[dev]"
python -m pic2smiles.cli env # confirms rdkit, torch and the accelerator
python scripts/capture_environment.py reports/environment.json
On a headless Linux node, RDKit's Cairo drawing also needs X libraries that a pip install does not bring:
Without them, from rdkit.Chem.Draw import rdMolDraw2D fails with
ImportError: libXrender.so.1.
MolScribe baseline (optional, Python 3.10, Linux):
conda create -y -n molscribe python=3.10 && conda activate molscribe
pip install torch==1.13.1 torchvision==0.14.1 --index-url https://download.pytorch.org/whl/cpu
pip install molscribe
pip install "opencv-python-headless==4.10.0.84" # not opencv-python: compute nodes lack libGL.so.1
pip install --no-deps "numpy==1.26.4" # rdkit 2022.09 predates numpy 2
curl -sSL -o swin_base_char_aux_1m.pth \
https://huggingface.co/yujieq/MolScribe/resolve/main/swin_base_char_aux_1m.pth
This environment does not work on macOS 26+; see docs/BASELINES.md.
1. Molecule corpus¶
curl -sSL -o chembl_37_chemreps.txt.gz \
https://ftp.ebi.ac.uk/pub/databases/chembl/ChEMBLdb/latest/chembl_37_chemreps.txt.gz
zcat chembl_37_chemreps.txt.gz | tail -n +2 | cut -f2 > chembl_37_smiles.txt
# Hand-drawn evaluation molecules, to be excluded from training
cut -f2 DECIMER_HDM_Dataset_SMILES.tsv | tail -n +2 > decimer_holdout.smi
python -m pic2smiles.cli corpus \
--source chembl_37_smiles.txt \
--output data/corpus/chembl \
--holdout decimer_holdout.smi \
--min-heavy-atoms 4 --max-heavy-atoms 48 \
--val-size 5000 --test-size 5000 --min-token-count 20 \
--seed 20260819
Writes train.smi, val.smi, test.smi, vocab.json and corpus_meta.json.
The metadata file records the curation counts and how many molecules the leakage
guard removed; the report embeds it verbatim as the declared scope.
2. Frozen evaluation sets¶
Evaluation images are written once and reused, so runs stay comparable.
# Rendered ChEMBL test molecules, three difficulty tiers
python -m pic2smiles.cli eval-set \
--smiles data/corpus/chembl/test.smi \
--output data/eval/chembl_test_linux \
--count 1000 --image-size 384 --seed 20260819
# Memorization-floor gallery: training molecules, clean style
head -3000 data/corpus/chembl/train.smi > gallery.smi
python -m pic2smiles.cli eval-set --smiles gallery.smi \
--output data/eval/gallery_train --count 3000 --styles clean --image-size 384
Build the eval set on the same host that trains, then copy the images to
wherever they are scored. Fonts differ between operating systems, and a set
rendered on another machine measures font generalization rather than accuracy
(see D14 in docs/DECISIONS.md). To measure that axis deliberately, build a
second set on the other host and score both.
The hand-drawn set comes from Zenodo 10.5281/zenodo.6456306; convert its TSV
into the manifest schema with columns
sample_id,source,split,style,canonical_smiles,image_path.
3. Train¶
Size the run from a measurement, not a guess:
python -m pic2smiles.cli train --corpus data/corpus/chembl --output runs/probe \
--preset base --steps 300 --batch-size 64 --amp bf16 --device cuda \
--eval-every 0 --checkpoint-every 0
Read img/s from the log; wall clock is steps * batch_size / img_per_second.
Then launch the real run:
python -m pic2smiles.cli train \
--corpus data/corpus/chembl --output runs/molmini_base \
--preset base --steps 240000 --batch-size 64 \
--lr 3e-4 --warmup 2000 --num-workers 14 \
--amp bf16 --device cuda \
--log-every 100 --eval-every 5000 --eval-samples 384 --checkpoint-every 5000 \
--seed 20260819
Or, under Slurm: sbatch scripts/slurm/train_a100.sbatch <corpus> <output> 240000 64 base.
Watch enc_cos in the log. It is the cosine similarity between the encoder's
memory for two different images. It should sit well below 1. Values approaching
1.0 mean the encoder has stopped distinguishing images and the decoder is
learning an unconditional SMILES prior — a failure a falling loss curve does not
reveal. The run aborts by itself above 0.995 after step 500.
Running from a git worktree¶
finalize.sh and the CLI resolve dataset paths under the repository root. A
fresh clone has a real data/ directory and needs nothing extra, but a git
worktree does not share the main checkout's ignored files, so link it:
4. Score and report¶
which is equivalent to, for each checkpoint and evaluation set:
python -m pic2smiles.cli benchmark \
--manifest data/eval/chembl_test_linux/manifest.csv --dataset chembl_test_linux \
--models molmini --checkpoint checkpoints/molmini_base.pt \
--batch-size 24 --output reports/results --run-id molmini_base_chembl_test_linux
python -m pic2smiles.cli benchmark \
--manifest data/eval/chembl_test_linux/manifest.csv --dataset chembl_test_linux \
--models oracle image_hash_nn --gallery data/eval/gallery_train/manifest.csv \
--output reports/results --run-id reference_chembl_test_linux
python -m pic2smiles.cli report --results reports/results \
--output reports/benchmark_report.md --scope reports/corpus_meta.json
oracle must score exactly 1.000. If it does not, the scoring path is broken
and no other number in the table is trustworthy.
Both best.pt and last.pt are scored. best.pt is chosen during training on
a 384-sample evaluation whose molecules differ between evaluations -- at a ~0.3
rate that carries roughly 2.3 points of standard error, so "best" can be a lucky
draw, and under a WSD schedule the genuinely best weights usually come last,
after the decay phase. The 3000-image frozen set is the estimate that can carry
a selection decision, so whichever checkpoint wins there is the one reported.
Comparing against MolScribe fairly¶
Only numbers produced by this harness on the same frozen manifest are
comparable. In particular, the [eval] lines in a training log are not
comparable to a benchmark table: they draw val-split molecules from an infinite
stream, render them fresh each time, and sample 384 of them. Different
molecules, different pixels, different sample each call. Use them to read a
trajectory, never to compare against another system.
5. Predict on your own image¶
Seeds and determinism¶
| what | seed | effect |
|---|---|---|
| corpus split | 20260819 |
which molecules land in train/val/test |
| frozen eval rendering | 20260819 |
the exact evaluation images |
| training data stream | --seed |
order and styling of training samples |
| in-training eval | 999983 (fixed) |
comparability across evals within a run |
Corpus curation and evaluation-set construction are deterministic: the same
inputs and seeds give byte-identical images, which
tests/test_data.py::test_eval_manifest_is_frozen_and_reproducible asserts.
Training is not bit-deterministic — GPU kernel scheduling is not — so a repeat
run lands near, not exactly on, the reported numbers.
Test suite¶
The tests worth knowing about are in tests/test_init_health.py: they assert
that encoder and decoder gradients are within 50x of each other at step 0, that
the encoder distinguishes real rendered depictions, and that zeroing the encoder
memory changes the decoder's logits. They fail on the pre-fix code and are the
guard against the failure documented as D18 in docs/DECISIONS.md.
Gotchas worth knowing¶
Do not edit finalize.sh while it is running. Bash reads a script lazily by
byte offset, so editing it mid-execution shifts the remaining bytes and the
shell misparses from wherever it had reached — the symptom is a syntax error on
a line that is perfectly valid. If you need to change it during a long run,
copy it first and run the copy.
Checkpoint sizes. A training checkpoint holds the live weights, the EMA
shadow and two AdamW moments — about four copies of the model, so 381 MB for the
23.8 M parameter base preset. pic2smiles export strips it to 95 MB, or 48 MB
with --half, and is verified to produce identical predictions.
Long rsync pulls. The checkpoints are the slow part of finalize.sh. Four
files at ~380 MB each takes a few minutes over a normal link.