File size: 3,030 Bytes
13e5482 1580d96 13e5482 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | ---
license: mit
library_name: pytorch
tags:
- biology
- protein-localization
- diffusion-model
- 3d
- cell-imaging
pipeline_tag: image-text-to-image
---
# DeepGPS-3D — baseline model weights
Pretrained weights for the **baseline DeepGPS-3D** model: a conditional 3D
denoising diffusion model that predicts a protein's 3D subcellular localization
volume from a matched nuclear 3D volume and the protein's ESM2 sequence
embedding.
- **Code / usage:** <https://github.com/royal-dargon/DeepGPS_3D>
- **File:** `best_model.pth` (~185 MB)
## Architecture
- **Diffusion:** DDPM, 1000 linear noise steps (β 1e-4 → 0.02).
- **Backbone:** conditional 3D U-Net, channel widths 32/64/128/256, two residual
blocks per level, 8-head attention at the bottleneck. Input is two spatial
channels (noisy protein volume + nucleus condition); output is the predicted
noise for a single protein channel.
- **Sequence conditioning:** precomputed ESM2-650M
(`facebook/esm2_t33_650M_UR50D`) 1,280-d embedding injected at every U-Net
scale.
- **Volume size:** 64 × 144 × 144 voxels.
- **Sampling:** classifier-free guided reverse diffusion (default guidance 3.0),
retaining nucleus and sequence conditions throughout.
## Training data
OpenCell confocal volumes, protein-disjoint split: 114,855 train images
(1,048 proteins) / 14,280 val (131) / 14,282 test (131). Trained with AdamW
(lr 1e-4), effective batch size 12 (batch 6 × grad-accum 2), up to 50 epochs,
seed 42.
## Usage
```bash
git clone https://github.com/royal-dargon/DeepGPS_3D
cd DeepGPS_3D
pip install -r requirements.txt
python scripts/download_weights.py --output weights/best_model.pth
python -m src.demo_infer \
--checkpoint weights/best_model.pth \
--nuclei example/WASL/WASL_nuclei.npy \
--embedding example/WASL/WASL_esm_650M_embedding.npz \
--reference example/WASL/WASL_reference.npz \
--output outputs/WASL_pred --device cuda
```
Or load directly:
```python
import torch
from src.model.diffusion import get_diffusion_model
from src.config import MODEL_CONFIG
model = get_diffusion_model(MODEL_CONFIG)
ckpt = torch.load("best_model.pth", map_location="cpu", weights_only=False)
state = ckpt["models"]["diffusion"] if isinstance(ckpt, dict) and "models" in ckpt else ckpt
model.load_state_dict(state.state_dict() if hasattr(state, "state_dict") else state)
model.eval()
# pred = model.sample_given_nuclei(nuclei[B,1,64,144,144], seq_emb[B,1280], guidance_scale=3.0)
```
## Intended use & limitations
Intended for research on protein subcellular localization and as the generative
backbone for the downstream mutation-localization and image-based PPI tasks in
the repository. The model reproduces coarse axial organization well and moderate
voxel/radial structure; it captures localization signal rather than exact
single-cell reconstruction. It is trained on one OpenCell split and reports
point estimates without confidence intervals. Not for clinical use.
## License
MIT. See the repository for details.
|