This project constructs a pipeline to get light models with minimal loss via distillation. First, we prune model structure of decoder and arquire light student models. Then, we do distillation training with teacher model dc-ae-f32c32-in-1.0 to reduce the gap as much as possible. During traing, we always freeze the encoder part and project out part of decoder. Furthermore, we involved tricks like AdamW optimizer, CosineAnnealingWarmRestarts Scheduler, Dynamic Loss Weight Adjustment Method, Batch Accumulation and Segment Training. Components of loss function include L1 distillation loss, L1 image loss, LPIPS loss, PatchGAN loss etc. After that, we choose FID, PSNR, SSIM and LPIPS as our evaluation matrice for image quality and MACs/inference time as indicators for speed. Finally, new light models outperform benchmark in PSNR and SSIM, the quality of generated images. Also, there is no significant difference in visualization between the generated image and teacher model.

The command of VAE model demonstration is in file "efficientvit/applications/dc_ae/scripts/demo_recons.sh", you may also use the command:

The demo picture will be stored in path "efficientvit/applications/dc_ae/reconstruction_results".

Code for modifying layers and pruning components is in "dc_de_modify_layer.py", and run the command:

There are three pruning methods. "Direct" means if I want to take 15 out of 30 parameters, the first 15 parameters will be taken directly. "gap" method will take 15 at intervals and "random" method will initialize to normally distributed random numbers.

Parameter pretrained_model are path for pretrained teacher model and save_path are for target model. In model save path, there're two other files, config.json for model_name which has been registered in ae_model_zoo.py and dc_ae.py and training_loss.txt for records during training.

Remember changing file/model path to the local path on your device. If you want to create new models by modifying layers, remember adding corresponding model info in "efficient/model/efficient/dc_ae.py and efficient/ae_model_zoo.py".

We have several versions of pruned models list below and you may download them from Google Drive links.

Normally, for all the training, we set 15 epochs and use 600 pictures randomly chosen from ImageNet for each epoch, costing about 75 mins on a single A6000 GPU. During all the training, we fix the encoder and project out layer of decoder.

There're two versions of pipeline, with GAN Loss and without GAN Loss but the main ideas of them are similar. For distillation part, there're three choices, before project_out layer, after TritonRMSNorm2d and after ReLU activation function. From ablation study, we choose to align the matrix after ReLU activation function using L1 Loss.

For the image generation loss part, we combine L1 loss, LPIPS loss and PatchGAN loss together with different weights.

For training techniques, we tried adding AdamW optimizer, CosineAnnealingWarmRestarts Scheduler, Dynamic Loss Weight Adjustment Method, Batch Accumulation and Segment Training, but only some of them seems to work. After that, we chose to use AdamW optimizer, CosineAnnealingWarmRestarts Scheduler and Dynamic Loss Weight Adjustment Method.

During experiments of parameter tuning, we found that distillation loss has much more importance than image generation loss. Even if the image evaluation matrices have reach a satisfactory result, the pictures are still not good enough due to high diatillation loss. Therefore, we give distillation loss a much higher weight during the first 10 epoch and let it gradually decrease to half of its original value according to the cosine law. The idea is to ensure the accuracy of distillation first! During the last 5 epoch, we focus more on image loss, and gradually increase their weights to double of their original value according to the cosine law.

For example, to train a student model based on w4-v3 pruned model, Loss = 100 * L1_Distillation + 1 * L1_Image + 0.1 * LPIPS_Image + 0.3 * PatchGAN_Image. And the training proportion of generation model(student model) and discrimination model is 300, which means we train student model for 300 samples and then train GAN discriminator model once, in case the discriminator learns too fast so that the student model will get lost and don't know how to optimize toward arquiring real images. The training command are as followed (the complete training command for three best models are in file train_distillation.sh):

You can measure the evaluation matrice (FID, PSNR, SSIM, LPIPS) of models through the command below, be aware to substitute the args "model" with your target model path.

During training, the generation pictures (ground truth, teacher model, student model) at the end of each epoch will be stored in "pic_save_dir" and the loss information will be save to "model_save_dir/training_losses.txt". So, we can better monitor the training process and analysis the problems.

Batch accumulation method use "accumulation_step" parameter to update parameter after certain batches to virtually increase batch size.

In Segement train method, we intend to train shallow layers/blocks first and deep ones later to resuce the training cost and improve efficiency. Unfortunately, these attempts seem to fail.

For exact and detailed meaning, type, default value etc. of each parameter, please refer to code or Appendix. If you are interested in more detailed training and debugging process, you may also refer to this Feishu Docs.

Generally the parameters and corresponding description are as followed: