UTRGAN
Model Introduction
UTRGAN is a generative and optimization framework for 5β² UTR design. It can generate candidate 5β² UTR sequences and predict and rank them according to gene expression, Mean Ribosome Load (MRL), and Translation Efficiency (TE).
The Hugging Face model package already contains the data and pretrained weights required for the basic workflow, so no additional model files need to be downloaded during inference.
Model Description
UTRGAN consists of multiple cooperating models:
- WGAN-GP Generator: Generates candidate 5β² UTR sequences of up to 128 nt from a 40-dimensional random latent vector;
- WGAN-GP Critic: Used during generative-model training;
- Xpresso: Predicts gene-expression-related scores;
- FramePool: Predicts Mean Ribosome Load (MRL);
- MTtrans: Predicts Translation Efficiency (TE);
- G4Boost: Performs G4-related classification and regression analyses.
WGAN, Xpresso, and FramePool use TensorFlow/Keras, while MTtrans uses PyTorch.
To maintain compatibility with the legacy H5 model files released by the original projects, Legacy Keras is used when loading the corresponding models.
Use Cases
| Use Case | Description |
|---|---|
| 5β² UTR candidate generation | Generate 5β² UTR candidate sequences in batches using the official pretrained WGAN. |
| MRL prediction and ranking | Use FramePool to calculate MRL scores and rank candidate sequences. |
| TE prediction and ranking | Use MTtrans to calculate TE scores and rank candidate sequences. |
| MRL/TE-directed optimization | Freeze the pretrained models and optimize the latent noise to improve the target score. |
| Gene-expression optimization | Combine UTRGAN with Xpresso to evaluate gene-expression-related properties of candidate sequences. |
| WGAN-GP training | Validate or retrain the generative model using the included UTRdb2 dataset. |
Usage
1. OneCode
You can use the OneCode online environment for an intelligent one-click AI4S programming experience:
Try OneCode for AI4S Programming
2. Manual Installation
Hardware Requirements
- CPU inference is supported;
- A OneScience-supported DCU environment is recommended for batch generation, ranking, and training;
- TensorFlow and PyTorch access the same DCU through their compatible accelerator interfaces;
- Full training time depends on the number of candidate sequences, batch size, and device performance.
Download the Model Package
Install the Hugging Face command-line tool and download the model repository:
python -m pip install -U huggingface_hub
hf download OneScience-Group/UTRGAN --local-dir ./UTRGAN
cd UTRGAN
Install the Runtime Environment
Create and activate a Python 3.11 environment, and then install the OneScience bioscience base environment:
conda create -n utrgan python=3.11 -y
conda activate utrgan
pip install onescience[bio-dcu]
Then install the additional or replacement dependencies specified in requirements.txt:
python -m pip install --no-deps -r requirements.txt
Note: This model adaptation uses DTK 26.04. The corresponding TensorFlow runtime has also been upgraded to TensorFlow 2.18.
Check the frameworks and available devices:
python - <<'PY'
import tensorflow as tf
import torch
print("TensorFlow:", tf.__version__)
print("TensorFlow devices:", tf.config.list_physical_devices("GPU"))
print("PyTorch:", torch.__version__)
print("HIP:", torch.version.hip)
print("DCU available:", torch.cuda.is_available())
PY
Weights and Data Preparation
The Hugging Face model package already contains the resources required for basic generation, prediction, ranking, and training validation:
| Resource | Location | Purpose |
|---|---|---|
| UTRdb2 | conf/data/utrdb2.csv |
WGAN-GP training data |
| Motif data | conf/data/motifs.csv |
Motif statistics and optimization analysis |
| WGAN Generator | weight/checkpoint_3000.h5 |
Generate candidate 5β² UTR sequences |
| FramePool | weight/utr_model_combined_residual_new.h5 |
MRL prediction |
| Xpresso | weight/humanMedian_trainepoch.11-0.426.h5 |
Gene-expression-related prediction |
| Xpresso | weight/GM12878_trainepoch.06-0.5062.h5 |
GM12878 expression-related prediction |
| Xpresso | weight/K562_trainepoch.11-0.4917.h5 |
K562 expression-related prediction |
| MTtrans | weight/mttrans/RL_hard_share_MTL/3R/schedule_MTL-model_best_cv1.pth |
TE prediction |
| G4Boost | weight/G4Boost_classifier.json |
G4 classification |
| G4Boost | weight/G4Boost_regressor.json |
G4 regression |
After applying the filtering and deduplication logic used by the official training scripts, UTRdb2 contains approximately 33,250 sequences with lengths between 65 and 128 nt.
The basic workflow does not require additional dataset downloads or precomputed features.
Optional Dependencies
requirements.txt includes XGBoost, ViennaRNA, logomaker, ruptures, and cliffs-delta for G4-related and downstream analysis scripts.
NUPACK is used only by optional MFE preprocessing scripts.
Because NUPACK 4 must be obtained and installed separately according to the licensing and installation requirements of its publisher, it is not automatically installed through standard PyPI dependencies in this model package.
NUPACK is not required for:
- Basic inference;
- MRL ranking;
- TE ranking;
- WGAN-GP training.
Quick Inference
Use the official pretrained WGAN to generate candidate sequences and rank them independently by MRL and TE.
Purpose: Generate 5β² UTR candidates in batches on DCU and output separate MRL and TE ranking results.
python scripts/predict.py \
--device dcu \
--device-id 0 \
--num-candidates 1024 \
--batch-size 128 \
--seed 33 \
--output-dir outputs/pretrained_batch_ranking
The results are saved to:
outputs/pretrained_batch_ranking/
βββ all_candidates_scores.csv
βββ ranked_by_mrl.csv
βββ ranked_by_te.csv
βββ generator_probabilities.npy
βββ summary.json
The output files contain:
all_candidates_scores.csv: all candidate sequences together with their MRL and TE prediction scores;ranked_by_mrl.csv: candidates ranked from highest to lowest MRL score;ranked_by_te.csv: candidates ranked from highest to lowest TE score;is_duplicate: indicates whether a generated sequence is duplicated;- MRL and TE are measured on different scales and their raw scores should not be directly added together.
MRL-Directed Optimization
Purpose: Freeze the WGAN and FramePool models and optimize the latent noise to increase the predicted MRL score of generated candidates.
python scripts/optimize_te_mrl.py \
-gpu 0 \
-task mrl \
-bs 64 \
-s 10 \
--output-dir outputs/optimization_mrl
During this process, the WGAN and FramePool model weights remain frozen. Only the latent noise is updated.
TE-Directed Optimization
Purpose: Freeze the WGAN and MTtrans models and optimize the latent noise to increase the predicted TE score of generated candidates.
python scripts/optimize_te_mrl.py \
-gpu 0 \
-task te \
-bs 64 \
-s 10 \
--output-dir outputs/optimization_te
During this process, the WGAN and MTtrans model weights remain frozen.
Therefore, this workflow is latent-space optimization rather than fine-tuning of the pretrained models.
Training
UTRGAN supports WGAN-GP training.
The trainable components in this workflow are:
- Generator;
- Critic.
Xpresso, FramePool, and MTtrans do not participate in this training entry point.
To run one epoch using the complete data-processing pipeline and the UTRdb2 dataset included in the repository:
python scripts/train.py \
-gpu 0 \
-bs 64 \
-lr 5 \
-mxl 128 \
-dim 40 \
--epochs 1 \
--output-dir outputs/train_1epoch
According to the upstream README, -lr 5 corresponds to a learning rate of:
1e-5
The adapted training entry point adds the --epochs argument and corrects the learning-rate expression in the upstream script so that it is consistent with the documented definition.
To run the full upstream training configuration, use the complete UTRdb2 training dataset for 4000 epochs:
python scripts/train.py \
-gpu 0 \
-bs 64 \
-lr 5 \
-mxl 128 \
-dim 40 \
--epochs 4000 \
--output-dir outputs/train_full
This retrains the WGAN-GP Generator and Critic.
Full training can require substantial runtime. Training progress should be evaluated according to available accelerator resources and training logs.
Generated checkpoints are saved in the specified output directory and do not overwrite the official pretrained weights under:
weight/
OneScience Official Resources
| Platform | OneScience Main Repository | Skills Repository |
|---|---|---|
| Gitee | OneScience | OneSkills |
| GitHub | OneScience | OneSkills |
Citation and License
- Upstream implementation: https://github.com/ciceklab/UTRGAN
- Xpresso: https://github.com/vagarwal87/Xpresso
- FramePool: https://github.com/Karollus/5UTR
- MTtrans: https://github.com/holab-hku/MTtrans
- The upstream UTRGAN project is distributed under the CC BY-NC-SA 2.0 license and is intended for academic and non-commercial use. Commercial use requires permission from the upstream authors.
- Third-party models, datasets, pretrained weights, and software components remain subject to their respective original licenses and terms of use.
- The DCU adaptation in this repository continues to use the same CC BY-NC-SA 2.0 license.