使用MTT实现数据集蒸馏:通过匹配训练轨迹生成小规模高质量合成数据集

mtt-distillation

数据集蒸馏简介

数据集蒸馏(Dataset Distillation)是一项具有挑战性的任务,旨在从大规模数据集中生成一个规模很小的合成数据集,使得在这个合成数据集上训练的模型能够达到与在原始大规模数据集上训练相当的性能。这项技术有望大大减少深度学习模型的训练时间和存储需求,对于资源受限的场景具有重要意义。

近年来,数据集蒸馏领域出现了多种方法,如梯度匹配、分布匹配等。本文介绍的MTT(Matching Training Trajectories)方法是由George Cazenavette等人在CVPR 2022会议上提出的一种新颖的数据集蒸馏技术,通过匹配训练轨迹来生成高质量的合成数据。

MTT方法原理

MTT的核心思想是优化合成数据,使得在合成数据上训练的网络能够产生与在真实数据上训练相似的参数轨迹。具体来说,MTT方法包含以下几个关键步骤:

在真实数据集上预训练多个"专家"网络,记录它们在训练过程中的参数轨迹。
初始化一组待优化的合成数据。
在合成数据上训练"学生"网络,同时计算学生网络参数与专家网络参数之间的距离。
通过反向传播,优化合成数据以最小化上述参数距离。
重复步骤3-4多次迭代,不断改进合成数据质量。

通过这种方式,MTT方法可以生成能够诱导相似网络训练动态的合成数据,从而在很小的数据规模下实现良好的模型性能。

MTT的实现与应用

环境配置

要使用MTT方法,首先需要配置相应的环境:

 git clone https://github.com/GeorgeCazenavette/mtt-distillation.git
 cd mtt-distillation

# 对于RTX 30XX及更新的GPU
 conda env create -f requirements_11_3.yaml

# 对于RTX 20XX及更早的GPU  
 conda env create -f requirements_10_2.yaml

conda activate distillation

生成专家轨迹

在进行蒸馏之前,需要先生成专家网络的训练轨迹:

 python buffer.py --dataset=CIFAR100 --model=ConvNet --train_epochs=50 --num_experts=100 --zca --buffer_path={path_to_buffer_storage} --data_path={path_to_dataset}

这个命令会在CIFAR-100数据集上训练100个ConvNet模型,每个模型训练50个epoch,并使用ZCA白化。

执行蒸馏过程

有了专家轨迹后,就可以开始蒸馏过程:

 python distill.py --dataset=CIFAR100 --ipc=1 --syn_steps=20 --expert_epochs=3 --max_start_epoch=20 --zca --lr_img=1000 --lr_lr=1e-05 --lr_teacher=0.01 --buffer_path={path_to_buffer_storage} --data_path={path_to_dataset}

这个命令会将CIFAR-100数据集蒸馏为每个类别只有1张图像的合成数据集。

在ImageNet子集上的应用

MTT方法还可以应用于ImageNet等大规模数据集的子集:

 python distill.py --dataset=ImageNet --subset=imagefruit --model=ConvNetD5 --ipc=1 --res=128 --syn_steps=20 --expert_epochs=2 --max_start_epoch=10 --lr_img=1000 --lr_lr=1e-06 --lr_teacher=0.01 --buffer_path={path_to_buffer_storage} --data_path={path_to_dataset}

这个命令会将ImageNet的"imagefruit"子集蒸馏为10张128x128分辨率的合成图像。

纹理合成:可穿戴的ImageNet

除了生成独立的合成图像,MTT方法还可以用于生成可平铺的纹理。这种纹理可以应用于衣物图案等需要连续性的场景。

要生成纹理,只需在distill.py命令中添加--texture标志:

 python distill.py --texture --dataset=ImageNet --subset=imagesquawk --model=ConvNetD5 --ipc=1 --res=256 --syn_steps=20 --expert_epochs=2 --max_start_epoch=10 --lr_img=1000 --lr_lr=1e-06 --lr_teacher=0.01 --buffer_path={path_to_buffer_storage} --data_path={path_to_dataset}

这个命令会将"imagesquawk"子集蒸馏为10个256x256分辨率的可平铺纹理。