Maximilian Müller and Matthias Hein
University of Tübingen and Tübingen AI Center
LoGex Workflow: First, an auxiliary classifier is trained on the long-tailed dataset. Then, StableDiffusion is finetuned via LoRa only on the tail samples. The auxiliary classifier and the finetuned diffusion model are used to generate synthetic tail images via guided diffusion. Those samples are then used to augment the tail classes of the real dataset, and a classifier is retrained on this augmented dataset. This improves the tail detection performance over existing methods and simultaneously achieves high classification accuracy on the head classes.
The repository is partly based on Gregory Holstes long-tailed chest x-ray classification repo and the AEDG repo (previously called DiG-IN).
For classifier training:
conda env create -f logex.yml
For synthetic data generation with AEDG:
cd AEDG
conda env create -f aedg.yml
The dataset used in this paper is a subset from Kriegsmann et al. You can download the full dataset from their website. The synthetic tail data generated with LoGex is available here. Downloading it directly allows you to jump to step 4 without fine-tuning stable diffusion and generating the tail samples yourself:
cd AEDG # data should be located in the AEDG folder
wget https://zenodo.org/records/13386825/files/syn_tail_images.tar.gz
tar -xvf syn_tail_images.tar.gz
Different choices are possible for the auxiliary classifier:
# normal cross-entropy
python src/main.py --data_dir /path/to/skincancer/data --out_dir output --dataset Skin_dataset_similar_PIL --loss ce --max_epochs 60 --patience 15 --batch_size 128 --lr 1e-4 --seed 1 --no_es --optim adamw --lrScheduler cosine --minLR 1e-5
# HOL loss
python src/main.py --data_dir /path/to/skincancer/data --out_dir output --dataset Skin_dataset_similar_PIL --loss ce --max_epochs 60 --patience 15 --batch_size 128 --lr 1e-4 --use_hol_new --lambda_hol 0.1 --lambda_hol_syn 0. --seed 1 --no_es --optim adamw --lrScheduler cosine --minLR 1e-5
# focal loss
python src/main.py --data_dir /path/to/skincancer/data --out_dir output --dataset Skin_dataset_similar_PIL --loss focal --max_epochs 60 --patience 15 --batch_size 128 --lr 1e-4 --seed 1 --no_es --optim adamw --lrScheduler cosine --minLR 1e-5
# ldam loss + cb reweighting
python src/main.py --data_dir /path/to/skincancer/data --out_dir output --dataset Skin_dataset_similar_PIL --loss ldam --max_epochs 60 --patience 15 --batch_size 128 --lr 1e-4 --seed 1 --no_es --optim adamw --lrScheduler cosine --minLR 1e-5 --rw_method cb
We used ldam+cb+drw for the numbers in the paper.
Finetune stable diffusion 1.4 via the diffusers library like explained here. Pretrained LoRa weights are also available here. Note that they differ slightly from the ones used for the experiments in the paper, so minor deviations are expected, but results should be similar.
cd AEDG
mkdir lora_weights # use path to this folder as model_path below
wget https://zenodo.org/records/13628864/files/pytorch_lora_weights.safetensors
Run the following commands from the AEDG folder (each command generates 100 samples from a tail class). The variables model_pathand weight_path need to be adjusted.
python src/cxr_classifier_conf.py gpu=0 loss=conf target_neurons=[674] target_classes=[0] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with dermis" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
python src/cxr_classifier_conf.py gpu=0 loss=conf target_neurons=[674] target_classes=[1] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with epidermis" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
python src/cxr_classifier_conf.py gpu=0 loss=conf target_neurons=[674] target_classes=[2] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with naevus" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
python src/cxr_classifier_conf.py gpu=0 loss=conf target_neurons=[674] target_classes=[4] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with subcutis" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
python src/cxr_classifier_conf.py gpu=0 loss=conf target_neurons=[674] target_classes=[5] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with sebaceousglands" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
python src/cxr_classifier_conf.py gpu=0 loss=conf target_neurons=[674] target_classes=[6] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with basal cell carcinoma" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
python src/cxr_classifier_conf.py gpu=4 loss=conf target_neurons=[674] target_classes=[8] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with skeletal muscle" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
python src/cxr_classifier_conf.py gpu=4 loss=conf target_neurons=[674] target_classes=[9] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with chrondal tissue" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
python src/cxr_classifier_conf.py gpu=4 loss=conf target_neurons=[674] target_classes=[10] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with sweat glands" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
python src/cxr_classifier_conf.py gpu=4 loss=conf target_neurons=[674] target_classes=[11] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with necrosis" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
python src/cxr_classifier_conf.py gpu=4 loss=conf target_neurons=[674] target_classes=[14] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with vessels" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
python src/cxr_classifier_conf.py gpu=4 loss=conf target_neurons=[674] target_classes=[15] sgd_steps=20 optim=adam sgd_stepsize=0.01 sgd_schedule=none latent_lr_factor=0.1 optimize_latents=True optimize_conditioning=True optimize_uncond=True loss_steps=3 loss_steps_schedule=linear per_timestep_conditioning=True per_timestep_uncond=True augmentation_num_cutouts=16 augmentation_noise_sd=0.005 augmentation_noise_schedule=const regularizers=[latent_background_l2,px_background_l2,px_background_lpips] regularizers_ws=[25.0,250.0,25.0] segmentation_squash=sqrt_0.3 prompt="A histopathological slide from a patient with elastosis" num_images=100 results_folder="skincancer_similar_pil_guided_class_specific_selectmaxconf_val" early_stopping_confidence=0.4 model_path=/path/to/lora/weights weight_path=/path/to/auxiliary/classifier/weights
CUDA_VISIBLE_DEVICES=3 python src/main.py --data_dir /path/to/skincancer/data --out_dir output --dataset Skin_dataset_similar_PIL_real_and_syn_targeted_maxconf_100 --loss ce --max_epochs 60 --patience 15 --batch_size 128 --lr 1e-4 --seed 1 --no_es --optim adamw --lrScheduler cosine --minLR 1e-5