模型稀疏加速训练
模型训练往往存在大量冗余计算,稀疏加速算法基于网络扩增训练的思想,结合参数继承方法,对模型中的部分算子实现权重的稀疏化。分别实现了宽度、深度层面的网络扩增算法,以应对不同场景,支撑模型训练阶段加速。
模型稀疏加速训练前需要参考环境准备完成开发环境部署、Python环境变量、所需框架及训练服务器环境变量配置。
目前支持对PyTorch框架下包括但不限于表1中的模型稀疏加速训练。
宽度扩增模型稀疏训练加速(以ResNet50模型为例)
- 用户需自行准备模型、训练脚本和数据集,本样例以PyTorch框架的ResNet50和数据集ImageNet为例。
- 编辑模型的训练脚本pytorch_resnet50_apex.py文件,导入稀疏加速训练的接口,新增代码参见如下示例。
from msmodelslim.pytorch import sparse
- (可选)新增如下代码,调整日志输出等级。启动训练任务后,将打屏显示调试的日志信息。
from msmodelslim import set_logger_level set_logger_level("info") #根据实际情况配置
- 当前训练脚本包含初始化模型、优化器optimizer和训练数据集,用户可以参考训练脚本pytorch_resnet50_apex.py自定义实现。在初始化模型、优化器optimizer和训练数据集后,新增如下代码调用稀疏化训练接口,将模型包装为稀疏化训练模型,请参考sparse_model_width进行配置。
epochs_each_stage = [10, 20, 30] #定义稀疏化每个阶段的epoch数量 model = sparse.sparse_model_width( model, optimizer, steps_per_epoch=len(train_loader), epochs_each_stage=epochs_each_stage )
其中epochs_each_stage = [10, 20, 30]表示稀疏化分3个阶段:
- 第1个阶段,从原模型裁剪为1/4的初始模型开始训练10个epoch。
- 第2个阶段将初始模型扩增2倍,训练20个epoch。
- 第3个阶段初始模型扩增为4倍,恢复为原模型大小训练,训练30个epoch。
若原训练脚本中已使用torch.nn.parallel.DistributedDataParallel封装了模型,稀疏加速接口sparse.sparse_model_width()的参数 model需为非ddp模式的模型。
- 启动训练任务,根据单卡或多卡调用不同的执行脚本,并指定data_path为数据集路径。
- 单卡训练时,执行命令启动训练任务。
bash ./test/train_full_1p.sh --data_path=./datasets/imagenet #请根据实际情况配置数据集路径
- 多卡训练时,执行命令启动训练任务。以下示例为8卡训练,请根据实际情况替换启动脚本。
bash ./test/train_full_8p.sh --data_path=./datasets/imagenet #请根据实际情况配置数据集路径
- 单卡训练时,执行命令启动训练任务。
深度扩增模型稀疏训练加速(以Swin-Transformer模型为例)
- 用户需自行准备模型、训练脚本和数据集,本样例以PyTorch框架的Swin-Transformer和数据集ImageNet为例。
- 编辑模型的训练脚本main.py文件,导入稀疏加速训练的接口,新增代码参见如下示例。
from msmodelslim.pytorch import sparse
- (可选)新增如下代码,调整日志输出等级。启动训练任务后,将打屏显示调试的日志信息。
from msmodelslim import set_logger_level set_logger_level("info") #根据实际情况配置
- 当前训练脚本包含初始化模型、优化器optimizer和训练数据集,用户可以参考训练脚本main.py自定义实现。在初始化模型、优化器optimizer和训练数据集后,新增如下代码,调用稀疏化训练接口,将模型包装为稀疏化训练模型,请参考sparse_model_depth进行配置。
epochs_each_stage = [10, 20, -1] #定义稀疏化每个阶段的epoch数量 model = sparse.sparse_model_depth(model, optimizer, steps_per_epoch=len(data_loader_train), epochs_each_stage=epochs_each_stage)
其中epochs_each_stage = [10, 20, -1]表示稀疏化分3个阶段:
- 第1个阶段,从原模型裁剪为1/4的初始模型开始训练10个epoch。
- 第2个阶段将初始模型扩增2倍,训练20个epoch。
- 第3个阶段epoch 数量“-1”表示训练直到总的 epoch 结束,初始模型扩增为4倍,恢复为原模型大小训练。
若原训练脚本中已使用torch.nn.parallel.DistributedDataParallel封装了模型,稀疏加速接口sparse.sparse_model_depth()的参数 model应为非ddp模式的模型。
- 启动训练任务,根据单卡或多卡调用不同的执行脚本,并指定data_path为数据集路径。
- 单卡训练时,执行命令启动训练任务。
bash ./test/train_full_1p.sh --data_path=./datasets/imagenet #请根据实际情况配置数据集路径
- 多卡训练时,执行命令启动训练任务。以下示例为8卡训练,请根据实际情况替换启动脚本。
bash ./test/train_full_8p.sh --data_path=./datasets/imagenet #请根据实际情况配置数据集路径
- 单卡训练时,执行命令启动训练任务。