sparse_model_width
功能说明
稀疏训练配置参数接口,判断当前稀疏化阶段,扩增模型与权重,以及调用重置 optimizer 接口,将用户提供的模型转化为稀疏化训练模型。
函数原型
sparse_model_width(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 = sparse.sparse_model_width(model, optimizer, steps_per_epoch=100, epochs_each_stage=[1, 2, 1])
父主题: 稀疏加速训练接口