YAML Metadata Warning:empty or missing yaml metadata in repo card
Check out the documentation for more information.
Universal Cell Embeddings
This repo is an UNOFFICIAL hugging implementation of the UCE (https://github.com/snap-stanford/UCE) model, designed for generate embeddings for single cell dataset. Based on the official repo of UCE, this repo provides scripts for pretraining, continue-training, and inference.
Installation
pip install -r requirements.txt
Dataset
A typical dataset folder containing .h5ad files looks like below:
├── dataset_uce # dataset root
│ ├── human # species
│ ├── human-sample1.h5ad
│ ├── human-sample2.h5ad
│ ├── ...
│ ├── mouse
│ ├── mouse-sample1.h5ad
│ ├── mouse-sample2.h5ad
│ ├── ...
│ ├── ...
Samples of the same species belong to the same species folder. The filenames of the species folder are used to map to the protein embedding model. See preprocessors.py for details.
More test samples can be found from the official dataset source
Train
We provide a simple example to show how to continue-train a 4 layer uce model. To train a model from scratch, one can simply skip the model loading part.
python train_hf.py --dataset_folder dataset/ \
--output_dir ./results/uce_ckpts/test-1 \
--dir test_output/ \
--model_loc model_files/4layer_model.torch \
--pretrained_model model-4 \
--logging_dir logs/exp1 \
--batch_size 4
Also, we can start the pretrain using deepspeed
deepspeed train_hf.py --dataset_folder dataset/ \
--output_dir ./results/test1 \
--dir test_output/ \
--model_loc model_files/4layer_model.torch \
--pretrained_model model-4 \
--logging_dir /logs \
--batch_size 4 \
--deepspeed_config ds_configs/zero2.json