DiT-Flow

Diffusion Transformer for Flow Cytometry

Pytorch implementation for the master thesis "Diffusion Transformers for MRD assessment in pediatric AML using Flow Cytometry" by Peter Buchegger.

Prerequisites

Tested on Windows 11 and Ubuntu 22.04 LTS with Python 3.10 and 3.13.

• Install Pytorch for your operating system and CUDA version
• pip install -r requirements.txt

Usage

Configure the hyperparameters and settings in a configuration file. A sample configuration file is given in config/default_config.yml.

Get dataset statistics

Before training is started the dataset preprocessing and filter can be investigated by invoking python get_dataset_stats.py -c myconfig.yml. It generates a file stats.txt containing the dataset statistics after preprocessing as it's done during training.

Training

All training and test outputs are saved in the out/[Trainer/TesterName]/[timestamp] where the checkpoints checkpoint[epoch].pth are generated and the best model is saved as best_model_checkpoint.pth.

Pretraining the FATE module of the dit-flow is started by python train_feature_encoder.py -c myconfig.yml.

After that the dit-flow model can be trained front-to-back using the pretrained model python train_dit.py -c myconfig.yml out/MaskedAEFATE/20260513-0843/best_model_checkpoint.pth.

Testing

For testing the MRD prediction issue python test_mrd_classification.py -c myconfig.yml out/DiTFlowTrainer/20260603-1955/best_model_checkpoint.pth.

Testing the dit-flow as FAB or Genotype classifier use these scripts

python test_fab_classification.py -c myconfig.yml out/DiTFlowTrainer/20260603-1955/best_model_checkpoint.pth

python test_geno_classification.py -c myconfig.yml out/DiTFlowTrainer/20260603-1955/best_model_checkpoint.pth

Generate sample

In order for the sample generation to be able to correctly decode the FATE latent the training has to be conducted with the -ff flag which keeps the pretrained FATE fixed while training the dit-flow diffusion part.

python train_dit.py -ff -c myconfig.yml out/MaskedAEFATE/20260513-0843/best_model_checkpoint.pth

The resulting model can generate new samples using the diffusion principle. Check the config parameters in the "generator_config" section.

python generate_sample.py -c myconfig.yml -o generated_sample.pkl out/DiTFlowTrainer/20260603-1955/best_model_checkpoint.pth out/MaskedAEFATE/20260513-0843/best_model_checkpoint.pth

Another script to compare the statistics of the generated latent before decoding through the FATE Decoder is provided to check whether the diffusion process does generate a latent with a standard deviation similar to that of the latent of a real world sample.

python compare_latent_statistics.py -c myconfig.yml out/DiTFlowTrainer/20260603-1955/best_model_checkpoint.pth out/MaskedAEFATE/20260513-0843/best_model_checkpoint.pth