深度学习框架-网络训练性能优化

184 阅读1分钟

背景:优化CPU后端常规网络训练性能,性能profiling采用关键步骤打点方式,性能优化时精度不变化。

1 网络脚本优化

1.1 消除训练后推理,复用训练的正向输出

class WithLossCell(nn.Cell):
  def construct(self, *data):
    ...
    out = self._backbone(*data)
    # return self._loss_fn(out, idx)
    return self._loss_fn(out, idx), ops.stop_gradient(out) ## ops.stop_gradient 截断梯度,屏蔽out对梯度影响

stop_gradient用法参考

1.2 batch计算代替for循环

# opt_ca_indexes = [valid_cell_idx.asnumpy()[ni] for ni in range(len(vertices))]
opt_ca_indexes = valid_cell_idx_np[:len(vertices)].tolist()

1.3 直接索引,不需要转numpy

# opt_ra_indexes = [valid_cell_idx.asnumpy()[ni] for ni in range(len(vertices))]
opt_ra_indexes = [valid_cell_idx_np[ni] for ni in range(len(vertices))]

1.4 优化代码,消除重复计算

2 高性能算子

分析发现matmul算子对接onednn、eigen算子库性能差异显著,arm机器eigen算子库性能高,x86机器环境onednn算子库性能高。

MindSpore提供Python API接口调用算子,在底层被调用算子根据实现逻辑分别接入CANN、CUDA、Eigen等框架。

mindspore三方算子库oneMKL、ondDNN路径:mindspore/ccsrc/plugin/device/cpu/kernel/mkldnn/

开源算子库:
oneMKL
oneDNN

eigen C++线性代数库
MindSpore入门:算子-浅入浅出