所谓Reduction(简约操作)其实就是说经过该操作之后,输入向量的维度会减小,例如,对输入向量沿着某个维度进行求和,先看一下python的代码:
import numpy as np
a = np.random.normal(size=(3, 4)).astype('float32')
print(a)
print(a.sum(axis=1))
打印结果为:
这里就是沿着第二个维度进行了求和操作,对应与具体实现如下:
import numpy as np
a = np.random.normal(size=(3, 4)).astype('float32')
print(a)
# print(a.sum(axis=1))
def sum_rows(a, b):
n = len(b)
for i in range(n):
b[i] = np.sum(a[i,:])
b = np.empty((3,), dtype='float32')
sum_rows(a, b)
print(b)
对应输出为:
过程相当直观,就是对第一维进行循环迭代,沿着第二维对数组进行求和,在numpy中采用了:符号实现了数组沿着第二维做的切片操作,在tvm中使用tvm.reduce_axis函数创建了tvm做简约操作的维度方向,需要指明起点和终点位置:
import tvm
n, m = tvm.var('n'), tvm.var('m')
A = tvm.placeholder((n, m), name='a')
j = tvm.reduce_axis((0, m), name='j')
B = tvm.compute((n, ), lambda i: tvm.sum(A[i, j], axis=j), name='b')
s = tvm.create_schedule(B.op)
print(tvm.lower(s, [A, B], simple_mode=True))
对应操作为:
~~~未完待续~~~
参考资料:
[1] https://github.com/d2l-ai/d2l-tvm
[2] http://tvm.d2l.ai.s3-website-us-west-2.amazonaws.com/