Awesome
基于深度学习的2D虚拟试衣
项目介绍
本项目主要面向第 $14$ 届全国服务外包创新创业比赛 $A16$ 赛道虚拟试衣赛题，采用 $2D$ 虚拟试衣技术依托于 $VITON$ 开源数据集训练 $DNN$ 网络并着重进行工程化落地应用；项目选用了前沿顶刊论文的 $PFAFN$ 模型，在此基础上对模型进行优化改进，实现了模型压缩和推理加速并使用 $OpenVINO$ 框架进行部署应用，出色地完成了赛题的要求。
项目开发环境
| 开发平台 | 版本 | 开发工具 | 版本 |
|---|
| Pycharm | 2022.3.2 | Visual Studio Code | 1.80.1 |
| Visual Studio | 17.5.5 | | |
| 开发环境 | 版本 | 开发环境 | 版本 |
|---|
| neural-compressor | 2.2.1 | nncf | 2.5.0 |
| numpy | 1.23.4 | onnx | 1.14.0 |
| opencv-python | 4.7.0.72 | onnxruntime | 1.15.1 |
| openvino | 2022.3.0 | pandas | 1.3.5 |
| pytorch-fid | 0.3.0 | rembg | 2.0.50 |
| pytorch | 2.0.0 | torch-pruning | 1.1.9 |
| intel-openmp | 2021.4.0 | | |
环境配置
git clone https://github.com/LZHMS/Virtual-Tryon.git
pip install -r requirements.txt
项目文件介绍
本项目主要分为模型训练和工程化落地两部分，因此仓库创建了两个分支 main 和 PruingQuantization。
main 分支是模型的推理部分，其中包括原始 Pytorch 模型、ONNX 模型、剪枝后模型、量化后模型的推理；
Img2Col模块用来对 corr_pure_torch模块做推理加速，模型训练中采用 corr_pure_torch模块而在推理阶段采用 Img2Col模块；- $afwm$ 与 $networks$ 分别是 $PFAFN$ 模型的衣服形变模块和图像生成模块
PruningQuantization 分支是模型工程化落地部分，其中还包括模型训练部分和模型剪枝量化；
ModelTraining 是 $PFAFN$ 模型的训练部分，分为四个阶段，先对教师网络进行训练然后采用可调节的知识蒸馏训练学生网络；ModelPruningQuantization 是本项目主要的工程化落地部分，模型剪枝主要针对于 $Warp$ 模型，为降低模型精度损失我们采用模块化剪枝的策略，并加入模型微调，将模型分为若干模块分别进行剪枝；模型量化采用了多种量化技术以及多种量化工具，具体尝试了 $Nerual\ Compressor$ 的训练后静态量化、Pytorch 的训练后静态量化以及 Pytorch 的量化感知训练。
模型结构介绍
本项目基于 $PFAFN$ 模型重新设计各个网络模块，具体结构如下图所示：
项目工程化落地
为了满足赛题方的要求，本项目开展了工程化落地部分，主要分为两个部分，模型训练和模型剪枝量化。项目工程化部署总图如下所示：
实验结果：通道剪枝
| Metrics | GFLOPs | Para(M) | SIZE(MB) | Total SIZE(MB) | Compresion Ratio | FID | FID Loss |
|---|
| Original Module | 6.63 | 9.37 | 35.8 | 112.0 | 100.00% | 8.906 | 0.00% |
| Ratio=0.2 with FineTuning | 5.23 | 7.28 | 27.6 | 88.69 | 79.19% | 9.013 | 1.20% |
| Ratio=0.3 with FineTuning | 4.40 | 6.48 | 24.8 | 65.73 | 58.69% | 9.113 | 2.32% |
| Ratio=0.4 with FineTuning | 3.79 | 5.61 | 20.4 | 40.97 | 36.58% | 9.304 | 4.47% |
| Ratio=0.5 with FineTuning | 3.42 | 4.55 | 16.8 | 35.47 | 31.67% | 9.977 | 12.03% |
| Metrics | GFLOPs | Para(M) | SIZE(MB) | Total SIZE(MB) | Compresion Ratio | FID | FID Loss |
|---|
| Original Module | 21.93 | 43.90 | 167 | 167 | 100.00% | 8.906 | 0.00% |
| Ratio=0.2 with FineTuning | 16.54 | 35.02 | 112.3 | 112.3 | 67.25% | 9.212 | 3.44% |
| Ratio=0.25 with FineTuning | 15.45 | 31.93 | 94.39 | 94.39 | 56.52% | 9.405 | 5.60% |
| Ratio=0.3 with FineTuning | 13.90 | 29.89 | 80.25 | 80.25 | 48.05% | 9.679 | 8.68% |
| Ratio=0.35 with FineTuning | 12.78 | 27.31 | 73.49 | 73.49 | 44.01% | 9.835 | 10.43% |
| Ratio=0.4 with FineTuning | 11.20 | 26.12 | 68.52 | 68.52 | 41.03% | 10.527 | 18.20% |
| Model | Original Model | Sparsity | Pruned Model | FID | FPS |
|---|
| CWM | 112MB | 40% | 40.97MB | 9.504 | 2.92 |
| IGM | 167MB | 25% | 94.39MB | 9.504 | 2.92 |
实验结果：量化感知训练
| Optimization | CPU-FID | GPU-FID | Original Model | Quantized Model |
|---|
| Unquantized | 9.504 | 9.483 | 135.36MB | 135.36MB |
| Quantize CWM | 9.783 | 9.701 | 40.97MB | 10.85MB |
| Quantize IGM | 10.382 | 10.249 | 94.39MB | 24.10MB |
| Quantize CWM & IGM | 11.503 | 11.379 | 135.36MB | 34.95MB |
实验结果：img2col 优化加速
| Runtimes | CorrTorch(s) | Img2Col(s) | FPS | Acceleration Rate |
|---|
| n=1000 | 147.8491 | 94.7902 | 10.81 | 1.5598 |
| n=10000 | 1489.1325 | 927.4293 | 10.77 | 1.6057 |
| Average Time | 0.1488 | 0.029 | 10.79 | 1.6017 |
参考文献
- Y. Ge, Y. Song, R. Zhang, C. Ge, W. Liu, and P. Luo, "Parser-Free Virtual Try-on via Distilling Appearance Flows," arXiv preprint arXiv:2103.04559, 2021.
- Y. Cheng, D. Wang, P. Zhou and T. Zhang, "Model Compression and Acceleration for Deep
Neural Networks: The Principles, Progress, and Challenges," in IEEE Signal Processing Magazine,
vol. 35, no. 1, pp. 126-136, Jan. 2018, doi: 10.1109/MSP.2017.2765695.
- PyTorch Quantization Aware Training