---
title: Training & tuning
description: Train a DTI model by pairing a drug and a protein encoder, tune train_epoch, LR, batch_size and cls_hidden_dims, then save the checkpoint with its config.
---

[DeepPurpose](https://github.com/kexinhuang12345/DeepPurpose) supports 50+ neural architecture combinations from 15+ compound encoders and 7+ protein encoders. This guide covers training DTI prediction models.

## Basic training loop

```python
from DeepPurpose import DTI as models
from DeepPurpose import utils, dataset

# Load and process data
X_drugs, X_targets, y = dataset.load_process_DAVIS(
    path='./data',
    binary=False,
    convert_to_log=True,
    threshold=30
)

# Process data with chosen encoders
train, val, test = utils.data_process(
    X_drug=X_drugs,
    X_target=X_targets,
    y=y,
    drug_encoding="Morgan",
    target_encoding="CNN",
    split_method="random",
    frac=[0.7, 0.1, 0.2],
)

# Configure and generate model
config = utils.generate_config(
    drug_encoding="Morgan",
    target_encoding="CNN",
    cls_hidden_dims=[1024, 1024, 512],
    train_epoch=100,
    LR=0.001,
    batch_size=256,
    cnn_target_filters=[32, 64, 96],
    cnn_target_kernels=[4, 8, 12],
)

model = models.model_initialize(**config)
model.train(train, val, test)
model.save_model("outputs/model")
```

## Architecture options

DeepPurpose uses a modular architecture where drug and target encoders are combined with a classifier head:

```
Drug SMILES --> [Drug Encoder] --\
                                  +--> [Classifier MLP] --> Affinity Score
Protein Seq --> [Target Encoder]-/
```

Common encoder combinations:
- **Morgan + AAC**: Fast baseline, good for small datasets
- **CNN + CNN**: Balanced performance and speed
- **Transformer + Transformer**: Best for large datasets
- **MPNN + CNN**: Graph-based drug representation with sequence-based protein encoding

## Hyperparameter tips

Key hyperparameters to tune:

| Parameter | Default | Description |
|-----------|---------|-------------|
| `train_epoch` | 100 | Number of training epochs |
| `LR` | 0.001 | Learning rate |
| `batch_size` | 256 | Training batch size |
| `cls_hidden_dims` | [1024, 1024, 512] | Classifier MLP hidden layer sizes |
| `cnn_drug_filters` | [32, 64, 96] | CNN filter sizes for drug encoder |
| `cnn_target_filters` | [32, 64, 96] | CNN filter sizes for target encoder |
| `transformer_n_layer_drug` | 8 | Transformer layers for drug encoder |
| `transformer_n_layer_target` | 8 | Transformer layers for target encoder |

Tuning recommendations:
- Start with small epochs (10-20) and increase after verifying loss trends.
- Use `batch_size` to fill GPU memory without OOM; scale `LR` with batch size cautiously.
- For small datasets (<1000 samples), use simpler encoders (Morgan, AAC) and smaller hidden dims.
- For GNN encoders (MPNN, DGL_*), reduce batch size as they are more memory-intensive.

## Validation and metrics

DeepPurpose automatically evaluates on validation data during training:

```python
# Training automatically reports validation metrics
model.train(train, val, test)

# For manual evaluation
from DeepPurpose import utils
y_pred = model.predict(test)

# Regression metrics
from sklearn.metrics import mean_squared_error, r2_score
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

# Classification metrics (if binary=True)
from sklearn.metrics import roc_auc_score, average_precision_score
auc = roc_auc_score(y_test, y_pred)
aupr = average_precision_score(y_test, y_pred)
```

The training process logs:
- Training and validation loss per epoch
- MSE, Pearson correlation, Concordance Index (CI) for regression
- AUROC, AUPRC for classification tasks

## Checkpoints

- Save `model.pt`, encoder configs, and the command/args used.
- Keep the tokenizer/featurizer versions with the checkpoint for
  reproducibility.

## Resource guidance

- CPU-only runs: lower batch size (e.g., 64) and fewer workers.
- GPU runs: ensure CUDA toolkit matches torch; set `TORCH_CUDA_ARCH_LIST` if
  building from source.

## Multi-task and virtual screening

DeepPurpose also supports:

```python
# Virtual screening: predict all drug-target pairs
from DeepPurpose import virtual_screening
result = virtual_screening.virtual_screening(
    X_drugs, X_targets, model, drug_names, target_names
)

# Drug repurposing for a specific target
from DeepPurpose import oneliner
oneliner.repurpose(
    target='MKKFFDSRREQGGSGLGSGSSGGGGSTSGLGSGYIGRVFGIGRQQVTVDEVLAEGGFAIVFLVRTSNGMKCALKRMFVNNEHDLQVCKREIQIMRDLSGHKNIVGYIDSSINNVSSGDVWEVLILMDFCRGGQVVNLMNQRLQTGFTENEVLQIFCDTCEAVARLHQCKTPIIHRDLKVENILLHDRGHYVLCDFGSATNKFQNPQTEGVNAVEDEIKKYTTLSYRAPEMVNLYSGKIITTKADIWALGCLLYKLCYFTLPFGESQVAICDGNFTIPDNSRYSQDMHCLIRYMLEPDPDKRPDIYQVSYFSFKLLKKECPIPNVQNSPIPAKLPEPVKASEAAAKKTQPKARLTDPIPTTETSIAPRQRPKAGQTQPNPGILPIQPALTPRKRATVQPPPQAAGSSNQPGLLASVPQPKPQAPPSQPLPQTQAKQPQAPPTPQQTPSTQAQGLPAQAQATPQHQQQLFLKQQQQQQQPPPAQQQPAGTFYQQQQAQTQQFQAVHPATQKPAIAQFPVVSQGGSQQQLMQNFYQQQQQQQQQQQQQQLATALHQQQLMTQQAALQQKPTMAAGQQPQPQPAAAPQPAPAQEPAIQAPVRQQPKVQTTPPPAVQGQKVGSLTPPSSPKTQRAGHRRILSDVTHSAVFGVPASKSTQLLQAAAAEASLNKSKSATTTPSGSPRTSQQNVYNPSEGSTWNPFDDDNFSKLTAEELLNKDFAKLGEGKHPEKLGGSAESLIPGFQSTQGDAFATTSFSAGTAEKRKGGQTVDSGLPLLSVSDPFIPLQVPDAPEKLIEGLKSPDTSLLLPDLLPMTDPFGSTSDAVIEKADVAVESLIPGLEPPVPQRLPSQTESVTSNRTDSLTGEDSLLDCSLLSNPTTDLLEEFAPTAISAPVHKAAEDSNLISGFDVPEGSDKVAEDEFDPIPVLITKNPQGGHSRNSSGSSESSLPNLARSLLLVDQLIDL',
    target_name='JAK3',
    drug_names=drug_list,
    save_path='./repurposing_result'
)
```