MindSpore简要性能分析

发布时间 2023-09-12 16:57:48作者: DECHIN

技术背景

在之前的一篇博客中,我们介绍过MindInsight的安装与使用。不过在前面的文章中,我们主要介绍的是MindInsight与SummaryCollector的配合使用,更多的是用于对结果进行回溯。这篇文章我们简要的从性能分析的角度,来介绍一下MindInsight的一些使用方法。

MindInsight的安装与启动

这部分内容在前面的博客中已经介绍过一次,这里简单的重复一下相关的内容。安装我们还是推荐使用pip进行安装和管理:

$ python3 -m pip install mindinsight

启动的方法很简单,就是在指定的目录下运行:

mindinsight start

如果在terminal里面显示如下内容,则表示安装成功。

Web address: http://127.0.0.1:8080
service start state: success

使用Profiler分析算子性能

当我们构建好相关的网络之后,类似于CPU中的line_profiler,这里我们可以用MindSpore中所支持的Profiler来直接进行网络性能评估。这里的Profiler主要以算子为单位进行统计,最终会输出每一个算子的调用次数以及相关的占用时长。使用方法非常简单,就是在代码的开头写一句:profiler = ms.Profiler(start_profile=True),以及在结尾处写一句:profiler.analyse()即可。

下面这个案例是MindSponge的一个能量极小化的案例。MindSponge是一个基于MindSpore框架开发的分子动力学模拟框架,更多的介绍和相关信息可以参考MindSponge教程系列博客。简单来说,这里我们只是模拟几百个水分子的动力学演化过程。

import os
os.environ['GLOG_v']='4'
os.environ['MS_JIT_MODULES']='sponge'
import mindspore as ms
from mindspore import context
from mindspore.nn import Adam


if __name__ == "__main__":

    import sys
    sys.path.insert(0, '..')

    from sponge import Sponge, Molecule, ForceField, WithEnergyCell
    from sponge.callback import RunInfo
    context.set_context(mode=context.GRAPH_MODE, device_target='GPU', device_id=1)

    profiler = ms.Profiler(start_profile=True)
    system = Molecule(template='water.tip3p.yaml')
    system.set_pbc_box([0.4, 0.4, 0.4])
    system.repeat_box([5, 5, 5])

    potential = ForceField(system, parameters=['TIP3P'], use_pme=False)

    opt = Adam(system.trainable_params(), 1e-3)

    sim = WithEnergyCell(system, potential)
    mini = Sponge(sim, optimizer=opt)

    run_info = RunInfo(10)
    mini.run(200, callbacks=[run_info])

    profiler.analyse()

该模拟过程的输出如下所示:

[MindSPONGE] Started simulation at 2023-09-11 10:57:03
[MindSPONGE] Compilation Time: 2.49s
[MindSPONGE] Step: 0, E_pot: 11003.434, Time: 2494.60ms
[MindSPONGE] Step: 10, E_pot: 9931.012, Time: 55.70ms
[MindSPONGE] Step: 20, E_pot: 9860.436, Time: 51.15ms
[MindSPONGE] Step: 30, E_pot: 9833.985, Time: 50.74ms
[MindSPONGE] Step: 40, E_pot: 9820.092, Time: 52.73ms
[MindSPONGE] Step: 50, E_pot: 9805.144, Time: 48.95ms
[MindSPONGE] Step: 60, E_pot: 9781.4375, Time: 47.87ms
[MindSPONGE] Step: 70, E_pot: 9740.486, Time: 48.46ms
[MindSPONGE] Step: 80, E_pot: 9684.311, Time: 48.67ms
[MindSPONGE] Step: 90, E_pot: 9619.269, Time: 52.02ms
[MindSPONGE] Step: 100, E_pot: 9555.06, Time: 51.39ms
[MindSPONGE] Step: 110, E_pot: 9498.154, Time: 47.72ms
[MindSPONGE] Step: 120, E_pot: 9450.13, Time: 49.52ms
[MindSPONGE] Step: 130, E_pot: 9408.7705, Time: 48.77ms
[MindSPONGE] Step: 140, E_pot: 9370.615, Time: 50.83ms
[MindSPONGE] Step: 150, E_pot: 9332.237, Time: 48.89ms
[MindSPONGE] Step: 160, E_pot: 9290.162, Time: 49.57ms
[MindSPONGE] Step: 170, E_pot: 9240.263, Time: 52.04ms
[MindSPONGE] Step: 180, E_pot: 9177.134, Time: 51.12ms
[MindSPONGE] Step: 190, E_pot: 9094.976, Time: 53.40ms
[MindSPONGE] Finished simulation at 2023-09-11 10:57:15
[MindSPONGE] Simulation time: 11.91 seconds.
--------------------------------------------------------------------------------

运行结束后,会在路径下生成一个data/profiler/的目录,里面存放有我们所需的性能分析相关信息。但是需要注意的是,profiler本身也会占用很多的运行时间,所以使用profiler和不使用profiler的运行时间会有比较大的差别。我们如果只是希望对算法代码进行优化,只要保持在同一个条件(使用或不使用profiler)下进行分析即可。

MindInsight性能分析

如果在前面代码运行的路径下直接启动MindInsight,然后把http://127.0.0.1:8080复制到浏览器里面打开,就可以看到对应的性能数据,如下图所示:

除了图形化的数据展示之外,还可以在列表中逐项展开,去查看每一个Operator的占用时间:

在这个结果中我们发现,由于使用了太多的Cast,有可能导致整体代码运行速度低下。如此一来,我们就可以挨个去查找代码中所用到的Cast算子,然后有针对性的进行优化:

$ grep -n -r "Cast" ~/projects/gitee/dechin/mindsponge/sponge/
/home/dechin/projects/gitee/dechin/mindsponge/sponge/potential/energy/coulomb.py:399:        self.cast = ops.Cast()
/home/dechin/projects/gitee/dechin/mindsponge/sponge/potential/energy/coulomb.py:490:        self.cast = ops.Cast()
$ grep -n -r "F.cast" ~/projects/gitee/dechin/mindsponge/sponge/
/home/dechin/projects/gitee/dechin/mindsponge/sponge/partition/fullconnect.py:75:            fc_idx = nrange + F.cast(no_idx <= nrange, ms.int32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/partition/distance.py:142:        max_neighbours = ops.count_nonzero(F.cast(mask, ms.float16), -1, dtype=ms.float16) - 1
/home/dechin/projects/gitee/dechin/mindsponge/sponge/partition/distance.py:143:        return F.cast(ops.reduce_max(max_neighbours), ms.int32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/partition/distance.py:239:            distances = F.cast(distances, ms.float16)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/partition/distance.py:242:            distances = F.cast(distances, ms.float32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/partition/grids.py:324:        sorted_grid_idx, sort_arg = self.sort(F.cast(atom_grid_idx, ms.float16))
/home/dechin/projects/gitee/dechin/mindsponge/sponge/partition/grids.py:325:        sorted_grid_idx = F.cast(sorted_grid_idx, ms.int32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/partition/grids.py:356:        grid_neigh_atoms, _ = self.sort(F.cast(grid_neigh_atoms, ms.float16))
/home/dechin/projects/gitee/dechin/mindsponge/sponge/partition/grids.py:357:        grid_neigh_atoms = F.cast(grid_neigh_atoms, ms.int32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/partition/grids.py:360:        max_neighbours = F.cast(msnp.amax(F.cast(max_neighbours, ms.float32)), ms.int32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/metrics/metrics.py:455:        classes_w_tensor = F.cast(classes_w_t2, mstype.float32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/metrics/metrics.py:482:        classes_num = F.cast(classes_num, mstype.float32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/metrics/metrics.py:541:        target = F.cast(target, mstype.float32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/metrics/metrics.py:542:        probs = F.cast(prediction, mstype.float32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/function/operations.py:295:                n = func.keepdims_sum(F.cast(mask, ms.int32), -2)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/function/functions.py:618:    return F.cast(image, ms.int32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/function/functions.py:1409:        value = F.cast(value, dtype)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/function/functions.py:1441:        value = F.cast(value, dtype)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/system/residue/residue.py:262:        self.natom_tensor = msnp.sum(F.cast(self.atom_mask, ms.float32), -1, keepdims=True)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/system/residue/residue.py:572:            F.cast(self.atom_mask, ms.int32), -1, keepdims=True)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/system/molecule/molecule.py:633:        self.system_natom = msnp.sum(F.cast(self.atom_mask, ms.float32), -1, keepdims=True)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/sampling/bias/metad.py:191:        cutoff_bins = F.cast(cutoff_bins, ms.int32)
/home/dechin/projects/gitee/dechin/mindsponge/sponge/sampling/bias/metad.py:398:        return F.cast(nearest_grid, ms.int32)

这里优化的思路和过程我们暂时就不做展示了,本文主要介绍的是一个性能优化的思路:先用profiler定位到性能决速步,然后有针对性的进行优化

MindInsight查看计算图

在使用AI框架进行计算的时候,我们的各种算子会被编译成一张大的计算图,使用MindInsight就可以对这个计算图进行可视化。使用方法也很简单,在设置context的时候多加上两项配置即可:

context.set_context(mode=context.GRAPH_MODE, device_target='GPU', device_id=1,)
                    save_graphs=True, save_graphs_path='./graphs/')

需要留意的是,这个save_graphs_path一定要配置上,否则输出的一大堆文件在当前目录下,直接没眼看。接下来同样的运行代码,会在当前目录下生成一个graphs/文件夹。此时刷新一下MindInsight的页面,点击最上面的训练列表,这个时候就会看到有两个数据列:

其中graphs这个数据列就是计算图的内容,点进去以后的界面如下所示:

如果我们比较关注计算图的话,就可以点进计算图的界面:

这个计算图的界面是根据右边的目录展开而展开的,如果我们想关注某一个模块的细节,就可以点进目录树:

这样的计算图结构,便于大家对整体的性能进行调节和优化。

总结概要

当我们需要优化程序性能的时候,首先我们就需要了解程序的主要耗时模块在哪里,也就是通常所谓的决速步,或者瓶颈模块,这样就可以有针对性的去进行优化。在MindSpore相关的程序中,我们可以使用MindInsight这一强力的性能分析可视化工具来进行分析。该工具会给出每个算子的调用次数以及总耗时等参数,能够给性能优化带来不少重要的参考。

版权声明

本文首发链接为:https://www.cnblogs.com/dechinphy/p/optimize.html

作者ID:DechinPhy

更多原著文章:https://www.cnblogs.com/dechinphy/

请博主喝咖啡:https://www.cnblogs.com/dechinphy/gallery/image/379634.html