Ensemble 3D UNet Soup

Model Details

Finetuned From Model

The Ensemble 3D UNet Soup model was pretrained on simulated datasets before fine-tuning on annotated cryoET data.

Model Architecture

3D U-Net with varying depth and channel sizes across model variants.

Parameters

• Tiny UNet: ~1.5M
• Medium UNet: ~4.5M
• Big UNet: ~11.8M

Citation

Peck, A., et al., (2025) A Real-World Phantom Dataset to Spur Innovation in CryoET Data Annotation. Accepted for publication in Nature Methods. Preprint: DOI: 10.1101/2024.11.04.621686

Primary Contact Email

Sirapoab Chaikunsaeng sirapoab.c@gmail.com, Sahil Barnwal confusedstudent13@gmail.com, Naoki Hashimoto maxashnitchts@yahoo.co.jp, Sergio Alvarez da Silva Junior sasjsergioalvarezjunior@gmail.com

To submit feature requests or report issues with the model, please open an issue on the GitHub repository.

System requirements

• NVIDIA T4
• RAM: ≥32GB
• OS: Ubuntu 22.04 or later

Model Variants

Intended Use

Primary Use Cases

• Localization of protein complexes within tomograms (e.g., ribosomes, virus-like particles, apo-ferritin, thyroglobulin, β-galactosidase)
• Automated annotation of subcellular components
• Preprocessing step for downstream 3D biological structure analysis

Out-of-Scope or Unauthorized Use Cases

Do not use the model for the following purposes:

• Use violating applicable laws or rights.
• Use for clinical diagnosis or decisions without expert validation.
• Use outside research and academic contexts.
• Any use that is prohibited by the Apache 2.0 license.
• Any use that is prohibited by the Acceptable Use Policy.

Training Details

Training Data

The training data was provided by the Chan Zuckerberg Imaging Institute (CZII) for the Kaggle competition and included seven experimental runs with ground truth annotations for six particles and 27 simulated runs. Please refer to the paper describing the CZII CryoET Object Identification Kaggle competition for more details.

Training Procedure

The models were pre-trained using the available simulated tomograms, with training performed on (128, 128, 128) 3D patches. Models were subsequently fine-tuned on the primary experimental dataset. Some models were trained solely on denoised tomograms, while others incorporated CTF-deconvolved and IsoNet-corrected tomograms.

We used the AdamW optimizer with an initial learning rate of 1×10^-3 and DiceCELoss as the loss function. Learning rate scheduling was handled via PyTorch Lightning's ReduceLROnPlateau callback.

Training included data augmentations consisting of random flips along the X, Y, and Z axes only. Some models incorporated Exponential Moving Average (EMA) with a momentum factor of 0.01, which helped smooth the learning curves and improve generalization.

The final models were combined using a "model soup" approach. The figure below summarizes this process:

Training Code

Training Hyperparameters

fp16 mixed precision.

Data Sources

CZII - CryoET Object Identification Challenge deposition site for:

• Experimental training data
• Simulated training data

Performance Metrics

Metrics

The model was evaluated by calculating the F-beta metric with a beta value of 4. The F-beta metric with a beta value of 4 is used to prioritize recall over precision, heavily penalizing missed particles while being more lenient on false positives. A particle was considered "true" if it lies within a factor of 0.5 of the particle of interest's radius. There were five particles of interest, with three "easy" particles (ribosome, virus-like particles, and apo-ferritin) assigned a weight of 1 and two "hard" particles (thyroglobulin and β-galactosidase) assigned a weight of 2. See CZII - CryoET Object Identification Evaluation description for more details, including the notebook with the official metric code.

Evaluation Datasets

The evaluation datasets included public and private test datasets found in the CryoET Data Portal deposition site for the Kaggle competition. The public and private test datasets contain 121 and 364 experimental runs, respectively. See:

• Public test dataset (Dataset ID: DS-10445)
• Private test dataset (Dataset ID: DS-10446)

Evaluation Results

Biases, Risks, and Limitations

Potential Biases

• Model may reflect annotation inconsistencies
• Model may underperform on rare structures or unseen imaging settings

Risks

• Inaccurate localization may lead to inaccurate biological interpretations
• Limited generalization beyond Cryo-ET phantom data

Limitations

• Model is not designed for clinical use
• Training data has limited augmentation and diversity

Caveats and Recommendations

• Validate model outputs before downstream use
• Review ensemble benefits vs. runtime tradeoffs
• We are committed to advancing the responsible development and use of artificial intelligence. Please follow our Acceptable Use Policy when using the model.

Acknowledgements

Thanks to the Chan Zuckerberg Initiative for hosting the competition. Special thanks to the Kaggle community and developers of PyTorch Lightning and MONAI.