模型稀疏加速训练

深度学习的一次训练往往涉及到几万甚至上百万的迭代次数，其中存在着大量的计算冗余，本算法基于网络扩增训练的思想，结合参数继承方法，旨在通过增加网络的宽度或深度来提升模型的性能，同时尽量减少训练成本和时间，以应对不同场景。

模型稀疏加速训练前需要参考环境准备完成开发环境部署、Python环境变量、所需框架及训练服务器环境变量配置。

目前支持对PyTorch框架下包括但不限于表1中的模型稀疏加速训练。

表1 已验证模型列表
类型	模型
语音识别	espnet-conformer
图像分类	ResNet50
	vgg16
	swin_tiny
目标检测	Yolov5
目标检测	Yolox

宽度扩增模型稀疏训练加速

导入依赖包并配置昇腾NPU。

import os
import torch
import torch_npu
import apex
from torch import nn
from apex import amp
from ascend_utils.common.utils import count_parameters
...
device = torch.device("npu:{}".format(os.getenv('DEVICE_ID', 0)))
torch.npu.set_device(device)

定义模型，并初始化优化器。

若为用户模型，需自行准备模型、训练脚本和数据集。

model = nn.Sequential(   
  nn.Conv2d(3, 32, 1, 1, bias=False), 
  nn.Sequential(nn.Conv2d(32, 64, 1, 1, bias=False), nn.BatchNorm2d(64), nn.Conv2d(64, 32, 1, 1, bias=False)), 
  nn.Sequential(nn.Conv2d(32, 64, 1, 1, bias=False), nn.BatchNorm2d(64), nn.Conv2d(64, 32, 1, 1, bias=False)),
  nn.Sequential(nn.Conv2d(32, 64, 1, 1, bias=False), nn.BatchNorm2d(64), nn.Conv2d(64, 32, 1, 1, bias=False)),
  nn.Sequential(nn.Conv2d(32, 64, 1, 1, bias=False), nn.BatchNorm2d(64), nn.Conv2d(64, 32, 1, 1, bias=False)),   
  nn.AdaptiveAvgPool2d(1),   
  nn.Flatten(),   
  nn.Linear(32, 10, bias=False), 
).to(device)

optimizer = apex.optimizers.NpuFusedSGD(model.parameters(), lr=0.1)

导入稀疏加速训练的接口。
```
from msmodelslim.pytorch import sparse
```
（可选）新增如下代码，调整日志输出等级。启动训练任务后，将会打屏显示调试的日志信息。
```
from msmodelslim import set_logger_level
set_logger_level("info")        #根据实际情况配置
```
当前训练脚本包含初始化模型、优化器optimizer和训练数据集。
在初始化模型、优化器optimizer和训练数据集后，新增如下代码调用稀疏化训练接口，将模型包装为稀疏化训练模型，请参考sparse_model_width进行配置。
```
steps_per_epoch, epochs_each_stage = 10, [2, 3, 1]   #定义稀疏化每个阶段的epoch数量
oring_model_params = count_parameters(model)  # 10826
model, optimizer = apex.amp.initialize(model, optimizer, opt_level="O2", combine_grad=False)  
# 添加宽度稀疏化训练方式
model = sparse.sparse_model_width(   
  model, optimizer, steps_per_epoch=steps_per_epoch, epochs_each_stage=epochs_each_stage 
)
```
其中epochs_each_stage = [2, 3, 1]表示稀疏化分3个阶段：
- 第1个阶段，从原模型裁剪为1/4的初始模型开始训练2个epoch。
- 第2个阶段将初始模型扩增2倍，训练3个epoch。
- 第3个阶段初始模型扩增为4倍，恢复为原模型大小训练，训练1个epoch。
若原训练脚本中已使用torch.nn.parallel.DistributedDataParallel封装了模型，稀疏加速接口sparse.sparse_model_width()的参数 model需为非ddp模式的模型。

启动训练任务。

for _ in range(steps_per_epoch * sum(epochs_each_stage)):
  optimizer.zero_grad()
  output = model(torch.ones([1, 3, 32, 32]).npu())
  loss = torch.mean(output)
  with amp.scale_loss(loss, optimizer) as scaled_loss:
      scaled_loss.backward()
  optimizer.step()

深度扩增模型稀疏训练加速

导入依赖包并配置昇腾NPU。

import os
import torch
import torch_npu
import apex
from torch import nn
from apex import amp
from ascend_utils.common.utils import count_parameters
...
device = torch.device("npu:{}".format(os.getenv('DEVICE_ID', 0)))
torch.npu.set_device(device)

定义模型，并初始化优化器。

若为用户模型，需自行准备模型、训练脚本和数据集。

model = nn.Sequential(   
  nn.Conv2d(3, 32, 1, 1, bias=False), 
  nn.Sequential(nn.Conv2d(32, 64, 1, 1, bias=False), nn.BatchNorm2d(64), nn.Conv2d(64, 32, 1, 1, bias=False)), 
  nn.Sequential(nn.Conv2d(32, 64, 1, 1, bias=False), nn.BatchNorm2d(64), nn.Conv2d(64, 32, 1, 1, bias=False)),
  nn.Sequential(nn.Conv2d(32, 64, 1, 1, bias=False), nn.BatchNorm2d(64), nn.Conv2d(64, 32, 1, 1, bias=False)),
  nn.Sequential(nn.Conv2d(32, 64, 1, 1, bias=False), nn.BatchNorm2d(64), nn.Conv2d(64, 32, 1, 1, bias=False)),   
  nn.AdaptiveAvgPool2d(1),   
  nn.Flatten(),   
  nn.Linear(32, 10, bias=False), 
).to(device)

optimizer = apex.optimizers.NpuFusedSGD(model.parameters(), lr=0.1)

导入稀疏加速训练的接口。
```
from msmodelslim.pytorch import sparse
```
（可选）新增如下代码，调整日志输出等级。启动训练任务后，将会打屏显示调试的日志信息。
```
from msmodelslim import set_logger_level
set_logger_level("info")        #根据实际情况配置
```
当前训练脚本包含初始化模型、优化器optimizer和训练数据集。
在初始化模型、优化器optimizer和训练数据集后，新增如下代码调用稀疏化训练接口，将模型包装为稀疏化训练模型，请参考sparse_model_depth进行配置。
```
steps_per_epoch, epochs_each_stage = 10, [2, 3, 1]   #定义稀疏化每个阶段的epoch数量
oring_model_params = count_parameters(model)  # 10826
model, optimizer = apex.amp.initialize(model, optimizer, opt_level="O2", combine_grad=False)  
# 添加宽度稀疏化训练方式
model = sparse_model_depth(   
  model, optimizer, steps_per_epoch=steps_per_epoch, epochs_each_stage=epochs_each_stage 
)
```
其中epochs_each_stage = [2, 3, 1]表示稀疏化分3个阶段：
- 第1个阶段，从原模型裁剪为1/4的初始模型开始训练2个epoch。
- 第2个阶段将初始模型扩增2倍，训练3个epoch。
- 第3个阶段初始模型扩增为4倍，恢复为原模型大小训练，训练1个epoch。
若原训练脚本中已使用torch.nn.parallel.DistributedDataParallel封装了模型，稀疏加速接口sparse_model_depth()的参数 model需为非ddp模式的模型。

启动训练任务。

for _ in range(steps_per_epoch * sum(epochs_each_stage)):
  optimizer.zero_grad()
  output = model(torch.ones([1, 3, 32, 32]).npu())
  loss = torch.mean(output)
  with amp.scale_loss(loss, optimizer) as scaled_loss:
      scaled_loss.backward()
  optimizer.step()