基于深度学习的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$ 模型重新设计各个网络模块，具体结构如下图所示：

项目工程化落地

为了满足赛题方的要求，本项目开展了工程化落地部分，主要分为两个部分，模型训练和模型剪枝量化。项目工程化部署总图如下所示：

实验结果：通道剪枝

• Clothe Warp Module

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%

• Image Generation Module

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