解析深度学习框架的“延迟计算”原理

Posted 2021-04-20 AI星课堂

大多数深度学习框架（如TensorFlow、pytorch、mxnet）都会使用延迟执行来提升系统的性能。什么是延迟计算？

什么是延迟计算？

延迟执行的大体意思就是：指命令可以等到之后它的结果真正的需要的时候再去执行，不到最后一步，不去执行。

下面，通过代码来讲解这个道理：

1a = 1+2
2# some other codes
3print (a)

通过第一句对a进行了赋值，然后再执行一些其指令后打印a的结果。因为这里我们可能很久以后才会用到a的值，所以，我们可以把它的执行延迟到最后真正执行它的时候。

这样做的目的：执行之前，系统可以看到后面的一系列命令，从而就会有更多的机会来对程序进行优化。

假如说，a在被打印前，被重新赋值了，那么，我们就可以不需要真正执行a=1+2这条语句了，因为a的值已经不再是3了。

大多数深度学习框架（如TensorFlow、pytorch、mxnet）都会使用延迟执行来提升系统的性能。绝大多数情况下，我们在写代码的时候，不用刻意去关注它的存在，因为大部分的时候，我们是不知道有这个工作在背后运行的。但是，理解它的工作原理，确实能够帮助我们进行非常合理和有效的开发。

深度学习框架的延迟计算

以深度学习框架mxnet为例，我们把用户直接写代码进行交互的部分叫做“前端”，比如从课程一开始到现在，我们一直在使用python写各种各样的方法和模型。除了python以外，mxnet还支持Scala，R，C++等语言作为其前端功能。

不管使用什么样的前端语言，mxnet的程序执行主要还是在C++后端进行。前端仅仅负责把程序传给后端。后端有自己的线程来不断的收集任务，构造计算图，优化，并且执行。

这一部分，就是要给小伙伴们介绍下后端的一个优化：延迟执行。

来看下面的代码：

1a = nd.ones((1,2))
2b = nd.ones((1,2))
3c = a*b+2
4print (c)

其中，我们在前端调用了4条语句，他们被后端的线程分析依赖并构建成计算图的模式。

分析依赖就是图中的箭头所指，变量是什么类型的，变量之间的运算关系。

计算图，就是图中由节点和边所组成的一个网络。

延迟执行的过程是这样的：

• 前端执行前三个语句的时候，它仅仅是把任务放进了后端的线程队列里就返回了。

• 当真正需要打印最终的运算结果的时候，前端会等待后端线程把c的结果计算完毕

这样设计的好处：

• 我们的前端是使用python进行交互的，大家都知道python运行是很慢的，但是我们加入延迟执行后，它并没有让python在前端参与任何的运算，所以，python前端对于整个程序的影响将会很小。

• 我们只需要后端使用C++来高效的运行，就可以很大程度生提升我们程序的性能了。

下面，我们通过代码来看看延迟执行的效果。可以看到当y=nd.dot(x,x)的时候，并没有等待它真正的被计算完。

1start = time()
2x = nd.random.uniform(shape = (2000,2000))
3y = nd.dot(x,x)
4print ('workloads are queued:\t%.4f sec'%(time()-start))
5print (y)
6print ('workloads are finished:\t%.4f sec'%(time()-start))

延迟执行对用户来说是透明的，很多时候，虽然我们写了很多代码，但是并不知道有这么一回事。除非我们需要打印或者保存结果外，我们根本是不关心目前是不是结果已经在内存中存放并且计算好了。

事实上，在我们写代码的过程中，只要把变量存在了ndarray里，而且使用了mxnet听的运算子。那么后端将默认我们使用了延迟执行这个选项，从而最大程度上的获取最大的性能。

延迟计算的好处

下面的例子中，我们对y进行了1000次的赋值，如果每次都需要等待y的值被计算出来，然后在进行下一次的迭代，那将会很慢。

如果我们使用了延迟执行，那么系统中就会忽略掉一部分的内容，从而加速运算。

 1x = nd.random.uniform(shape=(2000,2000))
 2start = time()
 3for i in range(1000):
 4    y = x+1
 5    y.wait_to_read()
 6print ('No lazy evaluation: %.4f sec'%(time()-start))
 7
 8
 9start = time()
10for i in range(1000):
11    y = x+1
12nd.waitall()
13print ('with evaluation: %.4f sec'%(time()-start))

上图左侧的依赖图深度为 7，右侧的依赖图深度为 3。计算产生相同的结果，但右侧的依赖图执行得更快。

延迟计算存在的问题

在延迟执行里，只要最终结果是一致的，系统可能使用跟代码不一样的顺序来执行，例如，我们写了如下代码：

1a = 1
2b = 2
3c = a+b

由于a=1和b=2之间没有任何依赖。所以说，我们使用b=2和a=1的顺序，也是能达到最终c=3的结果的。虽然对于最终的结果没有什么影响，但是，这样会导致内存使用的变化。

我们来列举几个在训练和预测中常见的现象。一般，对于每个批量我们都会评测一下，例如计算损失和精度，其中都会使用到asnumpy或者asscalar。这样，我们就能仅仅将一个批量的任务放进后端来执行。但是，如果我们去掉这些同步函数的话，就会使得大批量数据一次性全部放进系统进行执行，从而可能导致系统占用太多的资源而造成内存浪费。

下面，通过代码来讲讲上述延迟执行所带来的影响。

 1from mxnet import gluon
 2from mxnet.gluon import nn
 3from mxnet import ndarray as nd
 4from time import time
 5from mxnet import sym
 6import os
 7import subprocess
 8
 9def get_data():
10    start = time()
11    batch_size = 1024
12    for i in range(60):
13        if i%10==0 and i != 0:
14            print ('batch %d, time %f sec'%(i,time()-start))
15        x = nd.ones((batch_size,1024))
16        y = nd.ones((batch_size,))
17        yield x,y
18
19net = nn.Sequential()
20with net.name_scope():
21    net.add(
22        nn.Dense(1024, activation='relu'),
23        nn.Dense(1024, activation='relu'),
24        nn.Dense(1),
25    )
26net.initialize()
27trainer = gluon.Trainer(net.collect_params(), 'sgd', {})
28loss = gluon.loss.L2Loss()
29
30def get_mem():
31    res = subprocess.check_output(['ps','u','-p',str(os.getpid())])
32    return int(str(res).split()[15])/1e3
33
34for x,y in get_data():
35    break
36loss(y,net(x)).wait_to_read()
37
38mem = get_mem()
39for x,y in get_data():
40    loss(y,net(x)).wait_to_read()
41nd.waitall()
42print('Increased memory %f MB'%(get_mem()-mem))

结果如下：

1batch 10, time 0.213873 sec
2batch 20, time 0.436659 sec
3batch 30, time 0.663841 sec
4batch 40, time 0.883148 sec
5batch 50, time 1.115797 sec
6Increased memory 0.148000 MB

假设，此时，我们不适用wait_to_read()。那么前端就会将所有的批量一次性的全部添加进后端。可以看到，每个批量的数据都会在很短的时间内生成，同时在接下来的数秒内，我们看到了内存的显著增长情况。

1mem = get_mem()
2for x,y in get_data():
3    loss(y,net(x))
4nd.waitall()
5print('Increased memory %f MB'%(get_mem()-mem))

结果如下：

1batch 10, time 0.006374 sec
2batch 20, time 0.019963 sec
3batch 30, time 0.025409 sec
4batch 40, time 0.041614 sec
5batch 50, time 0.050791 sec
6Increased memory 55.120000 MB

总结，延迟执行使得系统有更多的内存空间来做性能的优化，但我们推荐每个批量里至少有一个同步函数，例如对损失函数进行评估，来避免将过多任务同时丢进后端系统。

解除延迟计算的方法

如何立即获取结果？除了前面介绍了的print功能外，我们还有别的方法可以让前端线程等待直到结果完成。

我们可以使用nd.ndarray.wait_to_read()等待直到特定结果完成，或者使用nd.waitall()等待所有前面的结果完成。

1start = time()
2y = nd.dot(x,x)
3y.asnumpy()
4print ('%.4f'%(time()-start))
5
6start = time()
7y = nd.dot(x,x)
8y.norm().asscalar()
9print ('%.4f'%(time()-start))

【1】https://zhuanlan.zhihu.com/p/35202071