Modelpublic

DnaRnaProteins/unet-bilstm-cell-cycle-baseline

sourceHugging Facecc-by-4.0updated 5mo agoView on Hugging Face

0likes5downloads

Model Card

U-Net + BiLSTM Cell Cycle Baseline

CPU-trained baseline checkpoint for 3-class cell cycle state prediction on CTC Fluo-N2DH-GOWT1. Architecture: 4-level U-Net encoder (no decoder) global avg pool 2-layer bidirectional LSTM linear classifier head.

Architecture

Component	Shape / config
Input clip	(B, T=8, 1, 64, 64) for the CPU variant; (B, T=16, 1, 96, 96) for the GPU config
Encoder	U-Net 4-level, base_ch=16 (CPU) / 32 (GPU)
Pooling	Global avg over (H, W)
Recurrent	BiLSTM, 2 layers, hidden=128 (CPU) / 256 (GPU), dropout=0.3
Head	Linear(2*hidden 3) on last timestep
Params	~0.95M (CPU) / ~3M (GPU)

Training

—Dataset: Fluo-N2DH-GOWT1 seq 01 (train) / seq 02 (val)
—Loss: weighted cross-entropy [1.0, 15.0, 25.0] (interphase / pre-mitosis / mitosis)
—Optimizer: AdamW, lr 3e-4, wd 1e-4, cosine schedule, grad clip 1.0
—30 epochs CPU / 80 epochs GPU

Held-out metrics (CPU, seq 02)

Metric	Value
Macro-F1	0.505
F1 interphase	0.98
F1 pre-mitosis	0.43
F1 mitosis	0.11
Mitosis event Â±3fr precision	0.15
Mitosis event Â±3fr recall	0.33

Files

—best_baseline_cpu.pt â state_dict for the CPU configuration
—config.yaml â hyperparameters and architecture choices

Loading

python

import torch
from huggingface_hub import hf_hub_download
# Architecture lives in this repo's training code:
#   github.com/mascharkh/biomechanical-ai-systems
#   project/src/cell_cycle/models/unet_baseline.py
from cell_cycle.models.unet_baseline import UNetLSTMClassifier

ckpt = hf_hub_download(
    repo_id="DnaRnaProteins/unet-bilstm-cell-cycle-baseline",
    filename="best_baseline_cpu.pt",
)
model = UNetLSTMClassifier(num_classes=3, base_ch=16, lstm_hidden=128, lstm_layers=2, dropout=0.3)
model.load_state_dict(torch.load(ckpt, map_location="cpu", weights_only=True))
model.eval()