Xiang-zx-zx/xenium-senescence-checkpoints
0
Xenium Senescence Checkpoints
UNet (42M params) trained with OT-CFM for cell trajectory generation.
Model Architecture
- UNet_CCDM_LatentCond: dim=64, dim_mults=[1,2,2,4,8], channels=1
- Input: 64×64 grayscale cell image
- Conditioning: senescence score (sinusoidal) + VAE latent (128-dim) + cell type embedding
- Output: velocity field v_pred (1, 64, 64)
Checkpoints
v8_{tissue}/model-50000.pt # V8 baseline (50K steps, constant LR 5e-5)
c25_20k/model-{step}.pt # C25 (cosine LR 2e-4, contrastive, 20K steps)Checkpoint Format
ckpt = torch.load("model-50000.pt")
# Keys: 'model' (state_dict), 'ema' (EMA state_dict), 'step' (int)Loading
from unet_ccdm_latent_cond import UNet_CCDM_LatentCond
unet = UNet_CCDM_LatentCond(
dim=64, dim_mults=[1,2,2,4,8], cond_emb_dim_init=128,
channels=1, num_cell_types=10, dropout=0.0
)
ckpt = torch.load("model-50000.pt", map_location="cpu")
unet.load_state_dict(ckpt['model'])Training Config (C25 — best FID)
accelerate launch --multi_gpu --num_processes 8 train_flow_matching.py \
--image_start --batch_size 512 --train_steps 20000 \
--lr 2e-4 --lr_schedule cosine --warmup_steps 2000 --lr_min 1e-6 \
--weight_decay 0.01 --contrastive_weight 0.1 \
--latent_drop_prob 0.5 --noise_prob 0.2 --noise_level 0.2 \
--num_cell_types 10 --auto_resume