sparse_model_depth

功能说明

稀疏训练配置参数接口，判断当前稀疏化阶段，裁剪或扩增模型与权重，以及调用重置 optimizer 接口，将用户提供的模型转化为稀疏化训练模型。

函数原型

sparse_model_depth(model, optimizer, steps_per_epoch, epochs_each_stage)

参数说明

参数名	输入/返回值	含义	使用限制
model	输入	初始化后的原始模型。若原训练脚本中已使用torch.nn.parallel.DistributedDataParallel封装了模型，model需为非ddp模式的模型。	必选。数据类型：PyTorch 模型。
optimizer	输入	初始化后的优化器optimizer。	必选。数据类型：PyTorch优化器，torch.optim.Optimizer的实例。
steps_per_epoch	输入	每个epoch的迭代数量，用于判断当前stage。	必选。数据类型：int，需大于0。
epochs_each_stage	输入	每个稀疏化阶段的epoch数量。	必选。数据类型：list或者tuple。元素必须是int。长度大于2，且其中元素除最后一个需要为大于0的int值，最后一个元素可以为-1。

调用示例

from modelslim.pytorch import sparse_model_depth
model = sparse.sparse_model_depth(model, optimizer, steps_per_epoch=100, epochs_each_stage=[10, 20, -1])

父主题： 稀疏加速训练接口