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

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:

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.