Retrosynthesis

The official TorchDrug 0.2.1 retrosynthesis tutorial implements the G2Gs workflow:

1. identify reaction centers,
2. split products into synthons,
3. complete synthons into reactants,
4. combine both trained tasks for end-to-end prediction.

This is a single-step reactant prediction pipeline. Multi-step route search, commercial availability, conditions, yields, and cost optimization are not provided by tasks.Retrosynthesis.

Prepare synchronized datasets

import torch
from torchdrug import datasets

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,
)

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

The repeated seed is required to align the reaction and synthon splits.

• Reaction mode stores (reactants, product) pairs.
• Synthon mode stores (reactant, synthon) pairs.

Center identification

The official tutorial uses RGCN and three feature groups:

from torchdrug import core, models, tasks

reaction_model = models.RGCN(
    input_dim=reaction_dataset.node_feature_dim,
    hidden_dims=[256, 256, 256, 256, 256, 256],
    num_relation=reaction_dataset.num_bond_type,
    concat_hidden=True,
)
reaction_task = tasks.CenterIdentification(
    reaction_model,
    feature=("graph", "atom", "bond"),
)

reaction_optimizer = torch.optim.Adam(
    reaction_task.parameters(),
    lr=1e-3,
)
reaction_solver = core.Engine(
    reaction_task,
    reaction_train,
    reaction_valid,
    reaction_test,
    reaction_optimizer,
    batch_size=128,
)
reaction_solver.train(num_epoch=50)
reaction_solver.evaluate("valid")
reaction_solver.save("g2gs-reaction.pth")

CenterIdentification predicts reaction centers. Its predict_synthon(batch, k=...) method returns top-k records containing synthons, reaction centers, reaction metadata, and log likelihoods.

Synthon completion

The official tutorial again uses RGCN:

synthon_model = models.RGCN(
    input_dim=synthon_dataset.node_feature_dim,
    hidden_dims=[256, 256, 256, 256, 256, 256],
    num_relation=synthon_dataset.num_bond_type,
    concat_hidden=True,
)
synthon_task = tasks.SynthonCompletion(
    synthon_model,
    feature=("graph",),
)

synthon_optimizer = torch.optim.Adam(
    synthon_task.parameters(),
    lr=1e-3,
)
synthon_solver = core.Engine(
    synthon_task,
    synthon_train,
    synthon_valid,
    synthon_test,
    synthon_optimizer,
    batch_size=128,
)
synthon_solver.train(num_epoch=10)
synthon_solver.evaluate("valid")
synthon_solver.save("g2gs-synthon.pth")

Do not substitute a GIN constructor copied from another implementation unless you intentionally redesign and validate the model.

End-to-end task

Combine the tasks, not the raw models:

task = tasks.Retrosynthesis(
    reaction_task,
    synthon_task,
    center_topk=2,
    num_synthon_beam=5,
    max_prediction=10,
)

If neither subtask has been attached to an Engine, preprocess them manually before composition:

reaction_task.preprocess(reaction_train, None, None)
synthon_task.preprocess(synthon_train, None, None)

The Retrosynthesis constructor accepts:

• center_identification
• synthon_completion
• center_topk
• num_synthon_beam
• max_prediction
• top-k metrics

It does not accept model=, synthon_model=, or a raw GNN pair.

Checkpoint loading

The official workflow saves each subtask and loads checkpoints without optimizer state when composing the pipeline:

optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver = core.Engine(
    task,
    reaction_train,
    reaction_valid,
    reaction_test,
    optimizer,
    batch_size=32,
)
solver.load("g2gs-reaction.pth", load_optimizer=False)
solver.load("g2gs-synthon.pth", load_optimizer=False)
solver.evaluate("valid")

Keep model architectures, feature sets, and dataset metadata identical to the training run. Inspect missing or unexpected keys if adapting this pattern.

Prediction output

The end-to-end task returns packed reactant predictions and a count per input:

from torchdrug import data, utils

batch = data.graph_collate(reaction_valid[:4])
batch = utils.cuda(batch)
predictions, num_prediction = task.predict(batch)

top1_index = num_prediction.cumsum(0) - num_prediction
for index in top1_index.tolist():
    reactants = predictions[index].connected_components()[0]
    print(reactants.to_smiles())

Call utils.cuda only when the task/models are on CUDA. Keep the batch on CPU for CPU execution.

Evaluation

The end-to-end task supports top-k exact-match metrics such as top-1, top-3, top-5, and top-10. Also inspect:

• chemical validity,
• duplicate predictions,
• performance by reaction class,
• atom-map consistency,
• stereochemistry retention,
• uncertainty or score gaps.

Top-k exact match against USPTO50k is not proof that a reaction is practical.

Scope and safety

TorchDrug's tutorial predicts reactant connectivity. It does not directly predict:

• reagents, catalysts, solvent, temperature, or pressure,
• yield or selectivity,
• commercial availability,
• multi-step search trees,
• process safety or scale-up feasibility.

Treat outputs as model proposals requiring forward validation, literature precedent, and expert chemistry review.

Common failures

Reaction and synthon samples do not align

Reset the same seed immediately before each split().

Feature dimension mismatch

Use the dedicated center_identification and synthon_completion atom features and derive model dimensions from each corresponding dataset.

End-to-end constructor error

Pass reaction_task and synthon_task, not their models.

Uninitialized metadata

Construct each subtask's engine first or call preprocess() manually.

Source links

• Retrosynthesis tutorial
• Retrosynthesis tasks
• USPTO50k dataset