URL

• https://mlc.ai/zh/chapter_integration/index.html

使用 Builder 创建 IRModule

从张量表达式创建 TensorIR（主张量函数）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

from tvm import te
# 定义 TensorIR 输入
A = te.placeholder((128, 128), name="A", dtype="float32")
B = te.placeholder((128, 128), name="B", dtype="float32")
type(A)
# tvm.te.tensor.Tensor
A.shape
# [128, 128]
# 由张量表达式自动生成 TensorIR
def te_matmul(A: te.Tensor, B: te.Tensor) -> te.Tensor:
    assert A.shape[1] == B.shape[0]
    n = A.shape[0]
    m = B.shape[1]
    k = te.reduce_axis((0, A.shape[1]), name="k")
    # 由张量表达式自动生成 TensorIR
    # 调用格式是：te.compute(output_shape, lambda, TensorIR_name)
    return te.compute(
        (n, m), lambda i, j: te.sum(A[i, k] * B[k, j], axis=k), name="matmul"
    )
C = te_matmul(A, B)
# 打印自动生成的 TensorIR，函数输入即为 [A, B, C]
te.create_prim_func([A, B, C]).show()

• 输出（自动生成的主张量函数）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

# from tvm.script import tir as T
@T.prim_func
def func(
    A: T.Buffer[(128, 128), "float32"],
    B: T.Buffer[(128, 128), "float32"],
    matmul: T.Buffer[(128, 128), "float32"],
) -> None:
    # function attr dict
    T.func_attr({"global_symbol": "main", "tir.noalias": True})
    # body
    # with T.block("root")
    for i0, i1, i2 in T.grid(128, 128, 128):
        with T.block("matmul"):
            i, j, k = T.axis.remap("SSR", [i0, i1, i2])
            T.reads(A[i, k], B[k, j])
            T.writes(matmul[i, j])
            with T.init():
                matmul[i, j] = T.float32(0)
            matmul[i, j] = matmul[i, j] + A[i, k] * B[k, j]

使用 BlockBuilder 构造 IRModule

• 自动生成的主张量函数还需要 计算图抽象 来将计算图拼起来

1
2
3
4
5
6
7
8
9
10
11
12

A = relax.Var("A", (128, 128), relax.DynTensorType(2, "float32"))
B = relax.Var("B", (128, 128), relax.DynTensorType(2, "float32"))
# 使用 BlockBuilder 将多个张量函数拼接成一个 IRModule
bb = relax.BlockBuilder()
with bb.function("main"):
    with bb.dataflow():
        C = bb.emit_te(te_matmul, A, B)
        D = bb.emit_te(te_relu, C)
        R = bb.emit_output(D)
    bb.emit_func_output(R, params=[A, B])
MyModule = bb.get()
MyModule.show()

• 输出（IRModule）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

tvm.script.ir_module
class Module:
    @T.prim_func
    def te_matmul(rxplaceholder: T.Buffer[(128, 128), "float32"], rxplaceholder_1: T.Buffer[(128, 128), "float32"], matmul: T.Buffer[(128, 128), "float32"]) -> None:
        ...
    @T.prim_func
    def te_relu(rxplaceholder: T.Buffer[(128, 128), "float32"], relu: T.Buffer[(128, 128), "float32"]) -> None:
        ...
    @R.function
    def main(A: Tensor((128, 128), "float32"), B: Tensor((128, 128), "float32")) -> Tensor(None, "float32", ndim = 2):
        # block 0
        with R.dataflow():
            lv = R.call_tir(te_matmul, (A, B), (128, 128), dtype="float32")
            lv1 = R.call_tir(te_relu, (lv,), (128, 128), dtype="float32")
            gv: Tensor((128, 128), "float32") = lv1
            R.output(gv)
        return gv

• 使用 BlockBuilder 创建 IRModule 与直接创建 IRMoudle 的对比
  
• bb.emit_te 做了以下事情：
  • 为 A 和 B 创建一个输入 te.placeholder
  • 通过 te_matmul 函数运行它们
  • 调用 te.create_prim_func 来创建一个 TensorIR 函数
  • 通过 call_tir 生成对函数的调用

从 Pytorch 映射到 IRModule

Pytorch 模型

1
2
3
4
5
6
7
8
9
10
11

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.weight = nn.Parameter(torch.randn(128, 128))
    def forward(self, x):
        x = torch.matmul(x, self.weight)
        x = torch.relu(x)
        return x
model = MyModel()
# 生成 Pytorch 计算图
fx_module = fx.symbolic_trace(model)

构造计算图之间的映射变换

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

# pytorch module parameter to IRModule parameter
def map_param(param: nn.Parameter):
    ndim = len(param.data.shape)
    return relax.const(
        param.data.cpu().numpy(), relax.DynTensorType(ndim, "float32")
    )
# pytorch module attribute to IRModule attribute
def fetch_attr(fx_mod, target: str):
    """Helper function to fetch an attr"""
    target_atoms = target.split('.')
    attr_itr = fx_mod
    for i, atom in enumerate(target_atoms):
        if not hasattr(attr_itr, atom):
            raise RuntimeError(f"Node referenced nonexistant target {'.'.join(target_atoms[:i])}")
        attr_itr = getattr(attr_itr, atom)
    return attr_itr
def from_fx(fx_mod, input_shapes, call_function_map, call_module_map):
    input_index = 0
    node_map = {}
    named_modules = dict(fx_mod.named_modules())
    bb = relax.BlockBuilder()
    fn_inputs = []
    fn_output = None
    with bb.function("main"):
        with bb.dataflow():
            for node in fx_mod.graph.nodes:
                if node.op == "placeholder":
                    # create input placeholder
                    shape = input_shapes[input_index]
                    input_index += 1
                    input_var = relax.Var(
                        node.target, shape, relax.DynTensorType(len(shape), "float32")
                    )
                    fn_inputs.append(input_var)
                    node_map[node] = input_var
                elif node.op == "get_attr":
                    node_map[node] = map_param(fetch_attr(fx_mod, node.target))
                elif node.op == "call_function":
                    node_map[node] = call_function_map[node.target](bb, node_map, node)
                elif node.op == "call_module":
                    named_module = named_modules[node.target]
                    node_map[node] = call_module_map[type(named_module)](bb, node_map, node, named_module)
                elif node.op == "output":
                    output = node_map[node.args[0]]
                    assert fn_output is None
                    fn_output = bb.emit_output(output)
        # output and finalize the function
        bb.emit_func_output(output, fn_inputs)
    return bb.get()

映射 Pytorch Module 到 TensorIR

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

# TensorIR 映射变换
def map_matmul(bb, node_map, node: fx.Node):
    A = node_map[node.args[0]]
    B = node_map[node.args[1]]
    return bb.emit_te(te_matmul, A, B)
# TensorIR 映射变换
def map_relu(bb, node_map, node: fx.Node):
    A = node_map[node.args[0]]
    return bb.emit_te(te_relu, A)
MyModule = from_fx(
    fx_module,
    input_shapes = [(1, 128)],
    call_function_map = {
      torch.matmul: map_matmul,
      torch.relu: map_relu,
    },
    call_module_map={},
)
MyModule.show()

• 映射后的 IRModule

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

@tvm.script.ir_module
class Module:
    @T.prim_func
    def te_matmul(rxplaceholder: T.Buffer[(1, 128), "float32"], rxplaceholder_1: T.Buffer[(128, 128), "float32"], matmul: T.Buffer[(1, 128), "float32"]) -> None:
        ...
    @T.prim_func
    def te_relu(rxplaceholder: T.Buffer[(1, 128), "float32"], relu: T.Buffer[(1, 128), "float32"]) -> None:
        ...
    @R.function
    def main(x: Tensor((1, 128), "float32")) -> Tensor(None, "float32", ndim = 2):
        # block 0
        with R.dataflow():
            lv = R.call_tir(te_matmul, (x, meta[relay.Constant][0]), (1, 128), dtype="float32")
            lv1 = R.call_tir(te_relu, (lv,), (1, 128), dtype="float32")
            gv: Tensor((1, 128), "float32") = lv1
            R.output(gv)
        return lv1

或映射到 Pytorch Module 到 IRModule 更高层的算子

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

def map_nn_relu_op(bb, node_map, node, nn_mod):
    A = node_map[node.args[0]]
    return bb.emit(relax.op.relu(A))
def map_nn_linear_op(bb, node_map, node, nn_mod):
    x = node_map[node.args[0]]
    w = map_param(nn_mod.weight)
    if nn_mod.bias is not None:
        b = map_param(nn_mod.bias)
    y = bb.emit(relax.op.dense(x, w))
    return bb.emit(relax.op.add(y, b))
MLPModuleHighLevel = from_fx(
    fx.symbolic_trace(mlp_model),
    input_shapes = [(1, 784)],
    call_function_map={
    },
    call_module_map={
        torch.nn.Linear: map_nn_linear_op,
        torch.nn.ReLU: map_nn_relu_op,
    },
)
MLPModuleHighLevel.show()

• 输出

1
2
3
4
5
6
7
8
9
10
11
12
13
14

@tvm.script.ir_module
class Module:
    @R.function
    def main(x: Tensor((1, 784), "float32")) -> Tensor(None, "float32", ndim = 2):
        # block 0
        with R.dataflow():
            lv: Tensor((1, 128), "float32") = relax.nn.dense(x, meta[relay.Constant][0])
            lv1: Tensor((1, 128), "float32") = relax.add(lv, meta[relay.Constant][1])
            lv2: Tensor((1, 128), "float32") = relax.nn.relu(lv1)
            lv3: Tensor((1, 10), "float32") = relax.nn.dense(lv2, meta[relay.Constant][2])
            lv4: Tensor((1, 10), "float32") = relax.add(lv3, meta[relay.Constant][3])
            gv: Tensor((1, 10), "float32") = lv4
            R.output(gv)
        return lv4

总结

• 张量表达式 API 允许我们创建原始的 TensorIR 函数
• BlockBuilder API 通过 emit_te 和其他函数创建 IRModule
• 通过将模型转换为 IRModule，实现与现有的机器学习框架的整合

URL

• https://mlc.ai/zh/chapter_auto_program_optimization/index.html

自动程序优化的原因

• MLC 的本质是张量函数之间的转换，但我们不知道哪种转换是让模型运行更快的，所以需要使用自动程序优化，去自动搜索最有转换。

自动程序优化过程

end-to-end 构建模型

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

# IR_Module 使用自定义主张量函数和库张量函数
@tvm.script.ir_module
class MyModuleMixture:
    @T.prim_func
    def linear0(
        X: T.Buffer[(1, 784), "float32"],
        W: T.Buffer[(128, 784), "float32"],
        B: T.Buffer[(128,), "float32"],
        Z: T.Buffer[(1, 128), "float32"],
    ):
        T.func_attr({"global_symbol": "linear0", "tir.noalias": True})
        ...
    @R.function
    def main(
        x: Tensor((1, 784), "float32"),
        w0: Tensor((128, 784), "float32"),
        b0: Tensor((128,), "float32"),
        w1: Tensor((10, 128), "float32"),
        b1: Tensor((10,), "float32"),
    ):
        with R.dataflow():
            lv0 = R.call_tir(linear0, (x, w0, b0), (1, 128), dtype="float32")
            lv1 = R.call_tir("env.relu", (lv0,), (1, 128), dtype="float32")
            out = R.call_tir("env.linear", (lv1, w1, b1), (1, 10), dtype="float32")
            R.output(out)
        return out
# 注册库张量函数
@tvm.register_func("env.linear", override=True)
def torch_linear(
    x: tvm.nd.NDArray, w: tvm.nd.NDArray, b: tvm.nd.NDArray, out: tvm.nd.NDArray
):
    ...
# 注册库张量函数
@tvm.register_func("env.relu", override=True)
def lnumpy_relu(x: tvm.nd.NDArray, out: tvm.nd.NDArray):
    ...
# 绑定模型权重参数（nd_params 是模型权重），作用类似于 functools.partial()
MyModuleWithParams = relax.transform.BindParams("main", nd_params)(MyModuleMixture)
# IR_Module -> 可执行程序 -> 虚拟机执行器
ex = relax.vm.build(MyModuleWithParams, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
# 执行
nd_res = vm["main"](data_nd)
# 测速
ftimer = vm.module.time_evaluator("main", tvm.cpu(), number=100)

自动优化 linear0 主张量函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

# 调优 API 只接受一个带有一个 main 函数的 IRModule，所以需要将原始 IRModule 中的 linear0 转成新 IRModule 的 main 函数
mod_linear = tvm.IRModule.from_expr(MyModuleMixture["linear0"].with_attr("global_symbol", "main"))
# 打印新IRModule
IPython.display.HTML(code2html(mod_linear.script()))
# 打印输出
@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(
        X: T.Buffer[(1, 784), "float32"],
        W: T.Buffer[(128, 784), "float32"],
        B: T.Buffer[(128,), "float32"],
        Z: T.Buffer[(1, 128), "float32"],
    ):
        # 函数中内容是 MyModuleMixture.linear0
# 自动调优 API，input 是需要调优的 IRModule，output 是调优后的 schedule，schedule.mod 是调优后的 IRModule
sch_tuned_linear = ms.tune_tir(
    mod=mod_linear,                 # 待调优 IRModule
    target="llvm --num-cores=1",    # 调优目标
    config=ms.TuneConfig(           # 自动调优配置
      max_trials_global=64,
      num_trials_per_iter=64,
    ),
    work_dir="./tune_tmp",
    task_name="main",
)
# 将返回的 IRModule 中的 main 函数更新到原 IRModule 的 linear0 中
# 绑定参数
MyModuleWithParams2 = relax.transform.BindParams("main", nd_params)(MyModuleMixture)
# 获取调优后的 main 函数
new_func = sch_tuned_linear.mod["main"].with_attr("global_symbol", "linear0")
# 获取原 IRModule 的 linear0 张量函数
gv = MyModuleWithParams2.get_global_var("linear0")
# 更新调优后的 main 函数到原 IRModule 的 linear0 张量函数
MyModuleWithParams2.update_func(gv, new_func)
# 重新测速，速度变快

URL

• https://mlc.ai/zh/chapter_end_to_end/index.html#

机器学习编译的本质与关注点

• MLC 的本质：张量函数之间的转换
• MLC 的关注点：
  • 所以可能的张量函数抽象表达
  • 所有可能的张量函数转换

构造 IR_Module

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def relu0(X: T.Buffer[(1, 128), "float32"], Y: T.Buffer[(1, 128), "float32"]):
        # function attr dict
        T.func_attr({"global_symbol": "relu0", "tir.noalias": True})
        for i, j in T.grid(1, 128):
            with T.block("Y"):
                vi, vj = T.axis.remap("SS", [i, j])
                Y[vi, vj] = T.max(X[vi, vj], T.float32(0))
    @T.prim_func
    def linear0(
        X: T.Buffer[(1, 784), "float32"],
        W: T.Buffer[(128, 784), "float32"],
        B: T.Buffer[(128,), "float32"],
        Z: T.Buffer[(1, 128), "float32"],
    ):
        T.func_attr({"global_symbol": "linear0", "tir.noalias": True})
        Y = T.alloc_buffer((1, 128), "float32")
        for i, j, k in T.grid(1, 128, 784):
            with T.block("Y"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]
        for i, j in T.grid(1, 128):
            with T.block("Z"):
                vi, vj = T.axis.remap("SS", [i, j])
                Z[vi, vj] = Y[vi, vj] + B[vj]
    @T.prim_func
    def linear1(
        X: T.Buffer[(1, 128), "float32"],
        W: T.Buffer[(10, 128), "float32"],
        B: T.Buffer[(10,), "float32"],
        Z: T.Buffer[(1, 10), "float32"],
    ):
        T.func_attr({"global_symbol": "linear1", "tir.noalias": True})
        Y = T.alloc_buffer((1, 10), "float32")
        for i, j, k in T.grid(1, 10, 128):
            with T.block("Y"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]
        for i, j in T.grid(1, 10):
            with T.block("Z"):
                vi, vj = T.axis.remap("SS", [i, j])
                Z[vi, vj] = Y[vi, vj] + B[vj]
    @R.function
    def main(
        x: Tensor((1, 784), "float32"),
        w0: Tensor((128, 784), "float32"),
        b0: Tensor((128,), "float32"),
        w1: Tensor((10, 128), "float32"),
        b1: Tensor((10,), "float32"),
    ):
        with R.dataflow():
            lv0 = R.call_tir(linear0, (x, w0, b0), (1, 128), dtype="float32")
            lv1 = R.call_tir(relu0, (lv0,), (1, 128), dtype="float32")
            out = R.call_tir(linear1, (lv1, w1, b1), (1, 10), dtype="float32")
            R.output(out)
        return out

• @tvm.script.ir_module 装饰 IR_Module
• @T.prim_func 装饰 主张量函数
• @R.function 装饰 高层神经网络执行的抽象，（将整个 IR_Module 中的主张量函数串起来组成一个计算图）
• R.dataflow() 用于标记程序计算图区域
• R.call_tir(prim_func, inputs, output_shape, dtype) 分配输出内存并和输入一起输入主张量函数

构建并运行模型

1
2
3
4
5

ex = relax.vm.build(MyModule, target="llvm")
vm = relax.VirtualMachine(ex, tvm.cpu())
nd_res = vm["main"](
    data_nd, nd_params["w0"], nd_params["b0"], nd_params["w1"], nd_params["b1"]
)

• 可执行文件 = build(IR_Module)
• 虚拟机执行器 = 虚拟机(可执行文件)
• 运行结果 = 虚拟机执行器(模型输入)

使用现有库避免重复造轮子

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

@tvm.script.ir_module
class MyModuleWithExternCall:
    @R.function
    def main(
        x: Tensor((1, 784), "float32"),
        w0: Tensor((128, 784), "float32"),
        b0: Tensor((128,), "float32"),
        w1: Tensor((10, 128), "float32"),
        b1: Tensor((10,), "float32"),
    ):
        # block 0
        with R.dataflow():
            lv0 = R.call_tir("env.linear", (x, w0, b0), (1, 128), dtype="float32")
            lv1 = R.call_tir("env.relu", (lv0,), (1, 128), dtype="float32")
            out = R.call_tir("env.linear", (lv1, w1, b1), (1, 10), dtype="float32")
            R.output(out)
        return out

• 上图中的 env.linear 是库张量函数，同一个 IR_Module 中可使用库张量函数，也可使用自定义张量函数，也可以二者混用。

1
2
3

@tvm.register_func("env.linear", override=True)
def torch_linear(x: tvm.nd.NDArray, w: tvm.nd.NDArray, b: tvm.nd.NDArray, out: tvm.nd.NDArray):
    ...

• @tvm.register_func，注册库张量函数

绑定参数到 IR_Module

• 以上的所有 IR_Module 都是在调用时才传入 数据 和 权重参数，但 权重参数 是不变的，所以可以将 权重参数 提前绑定到 IR_Module 中，调用时只传入输入数据。

1

MyModuleWithParams = relax.transform.BindParams("main", nd_params)(MyModuleMixture)

• relax.transform.BindParams(计算图入口函数，模型参数)(IR_Module) 将模型参数绑定到 IR_Module 的计算图入口函数中，返回一个绑定好模型参数的 IR_Module。

总结

• 计算图抽象（一般表示 main 函数） 有助于将元张量函数拼接在一起以进行端到端执行。
• Relax 抽象的关键要素包括
  • call_tir 构造，将目标传递规范的元函数嵌入到计算图中
  • Dataflow block
• 计算图允许调用环境库函数和 TensorIR 函数。

三种成像相关坐标系

• 像素坐标系：一种 2D 坐标，以图片左上角为原点，横轴（宽度方向）向右为 x 轴正方向，纵轴（高度方向）向下为 y 轴正方向，单位是像素
• 相机坐标系：一种 3D 坐标，以相机光心为原点，垂直相机平面远离相机方向为 z 轴正方向，垂直于 z 轴且平行于相机平面，水平向右为 x 轴正方向，竖直向下为 y 轴正方向，单位是米
• 世界坐标系：一种 3D 坐标，一种人为定义的，且x, y, z 轴两两垂直的坐标系，单位是米

相机内参

• 相机内参可以实现像素坐标系与相机坐标系之间相互转换，通常使用一个 3 * 3 矩阵表示
  
  

• 根据小孔成像和相似三角形原理，可以得出相机坐标系与成像坐标系点的对应关系：$\frac{Z}{f}=\frac{X}{x}=\frac{Y}{y}$
  其中：$(X,Y,Z)$ 为相机坐标系下的点的坐标， $(x, y)$ 为投影到成像平面上的点的坐标， $f$ 表示焦距。

• 再根据成像坐标系到像素坐标系的对应关系：

  $$\begin{matrix} u=\alpha \cdot x + c_x \\ v = \beta \cdot y + c_y \end{matrix}$$

  其中：

  • $\alpha,\beta$ 分别表示 $x, y$ 方向上成像宽度到像素宽度的投影
  • 由于成像坐标系原点为成像中心，像素坐标系原点为像素左上角，所以需要加上原点的偏移， $c_x,c_y$ 分别表示 $x, y$ 方向上原点的偏移。

• 所以，相机坐标系 $(X,Y,Z)$ 与像素坐标系 $(u, v)$ 可通过相机内参相互转换：

  $$Z \begin{bmatrix} u \\ v \\ 1 \end{bmatrix} = \begin{bmatrix} f_x & 0 & c_x \\ 0 & f_y & c_y \\ 0 & 0 & 1 \end{bmatrix} \begin{bmatrix}X\\ Y\\ Z \end{bmatrix}$$

  其中： $f_x=\alpha \cdot f,\ f_y=\beta \cdot f,\ Z$ 表示相机坐标系下的深度

• $\begin{bmatrix} f_x & 0 & c_x \\ 0 & f_y & c_y \\ 0 & 0 & 1 \end{bmatrix}$ 被称为相机内参 $K$

相机外参

• 相机外参可以实现相机坐标系与世界坐标系之间相互转换（刚体变换），通常用一个 3 * 3 的旋转矩阵 $R$ 和一个 3 * 1 的平移矩阵 $T$ 表示：

  $$\begin{bmatrix} X_c\\ Y_c \\ Z_c \end{bmatrix}=\begin{bmatrix} R_{11} & R_{12} & R_{13} \\ R_{21} & R_{22} & R_{23} \\ R_{31} & R_{32} & R_{33} \end{bmatrix} \begin{bmatrix} X_w\\ Y_w \\ Z_w \end{bmatrix} + \begin{bmatrix} T_{1} \\ T_{2} \\ T_{3} \end{bmatrix}$$

  其中： $(X_c,Y_c,Z_c)$ 表示相机坐标系下的点， $(X_w,Y_w,Z_w)$ 表示世界坐标系下的点

• 齐次化之后，得到一个 4 * 4 的矩阵：

  $$\begin{bmatrix} X_c\\ Y_c \\ Z_c \\ 1\end{bmatrix}=\begin{bmatrix} R_{11} & R_{12} & R_{13} & T_1 \\ R_{21} & R_{22} & R_{23} & T_2 \\ R_{31} & R_{32} & R_{33} & T_3 \\ 0 & 0 & 0 & 1\end{bmatrix} \begin{bmatrix} X_w\\ Y_w \\ Z_w \\ 1\end{bmatrix}$$

总结

• $camera\ intrisic\ matrix \times camera\ coordinate\ system = depth\ in\ camera\ coordinate system \cdot image\ coordinate\ system$
• $camera\ extrisic\ matrix \times world\ coordinate\ system = camera\ coordinate\ system$
• $camera\ intrisic\ matrix \times camera\ extrisic\ matrix \times world\ coordinate\ system = depth\ in\ camera\ coordinate system \cdot image\ coordinate\ system$

URL

• paper: https://arxiv.org/pdf/2008.05711.pdf
• code: https://github.com/real-zhangzhe/lift-splat-shoot (fork from official implementation, adding comments and predictions)

TL;DR

• 本文是 BEV (bird eye view) 的开山之作，通过隐式 2D 深度估计和像素坐标到世界坐标转换，将多张（6张）车周环视图拼接得到一张鸟瞰图。
• 具体实现请看代码，代码中有非常详细的注释。

Algorithm

在 inference 过程中，下半部分的 6 张环视图为输入，上半部分的鸟瞰图为输出（地图和本算法无关）。

Dataset

• 本文使用自动驾驶数据集 nuScense
  • 官网：https://www.nuscenes.org/
  • 数据集介绍论文：https://arxiv.org/pdf/1903.11027.pdf
    
• 输入的 6 张图来自上图的 6 个绿色 camera
• 世界坐标系如图 IMU 所示原点定为 车后轴中心，x 轴正方向为车辆前进方向， y 轴正方向为面向车辆前进方向的左手边，z 轴正方向为竖直向上。

相关知识

• 相机内外参：参考 相机内参与外参
• 体素：
  • 立体像素
  • 在本文中，一个体素表示 x 方向上 0.5m，y 方向上 0.5m，z 方向上 20m 的立方体区域
  • 体素为鸟瞰视角中可区分的最小单元

算法细节

特别细节的看代码

1. 特征提取（Lift）

• 使用 EfficientNet 进行 2D 特征提取 + 隐式深度估计
  • 输入：shape = [4, 6, 3, 128, 352]，分别表示：[batch, cameras, channel, height, width]
  • 输出：shape = [24, 64, 41, 8, 22]，分别表示：[batch * cameras, features, depth, height, width]
    • 使用 64 维向量编码深度（不是直接预测深度，所以被称为隐式深度估计）
    • 深度从 4m ~ 45m，编码精度为 1m，所以有 41 种离散深度，相机坐标系下深度估计的目的是：从像素坐标转化为世界坐标
    • 长宽各下采样 16 倍，减小计算量

2. 像素坐标和相机坐标系下深度到世界坐标的映射（Splat）

• 使用如下参数将像素坐标和相机坐标系下深度映射到世界坐标
  • 相机内参
  • 相机外参
    • 旋转
    • 平移
  • 像素坐标系内变换参数（缩放 + 裁剪（平移））
    • 原图(900, 1600) -> 模型输入图(128, 352) -> 模型预测图(8, 22)
• 体素池化：将属于同一个体素的深度估计向量求和
• 输入：
  • 深度估计：shape = [24, 64, 41, 8, 22]
  • 相机内外参和缩放参数
• 输出：shape = [4, 64, 200, 200]
  • 200 * 200 个体素
    • X 方向上 [-50m, 50m) 0.5m 为一个 bin，200 个 bin
    • Z 方向上 [-50m, 50m) 0.5m 为一个 bin，200 个 bin
    • Y 方向不分 bin
  • 每个体素用 64 维向量编码
• 本质是：
  1. 构造一个 [24 * 41 * 8 * 22, 3] 的查找表，输入为 backbone 输出特征图的每一个 pixel，输出为这个 pixel 对应的世界坐标（这个查找表可由相机内外参和图像缩放系数计算得到）
  2. 将离散的世界坐标点合并，合并规则是属于同一个体素的坐标点则合并

3. 体素编码降维（Shoot）

• 输入：shape = [4, 64, 200, 200]
• 输出：shape = [4, 1, 200, 200]，BEV 图

4. 训练 loss

非常简单粗暴

• 将 GT bbox3d 同样映射到 BEV 空间 [4, 1, 200, 200]，然后做 pixel-wise loss（分割 loss）

Thought

• 代码中 像素坐标和相机坐标系下深度到世界坐标的映射 部分比较难懂，需要较强的相机成像原理 / 图像 2D 转 3D 背景才能看懂

URL

• paper: https://arxiv.org/pdf/2103.14030.pdf
• code: https://github.com/microsoft/Swin-Transformer

TL;DR

• 本文提出一种新型 Transformer 结构使用 Window Multi-head Self-attention (W-MSA) 和 Shifted Window Multi-head Self-attention (SW-MSA) 结构替代原始 Transformer 使用的 Multi-head Self-attention (MSA) 结构。
  • 大大节省了原始 Transformer 的计算复杂度。
  • 在视觉任务中吊打了一众 CNN 和 Transformer。

Algorithm

总体结构

• swin-Transformer 总体按照类似 CNN 的层次化构建方式构建网络结构，分为 4 个 stage，每个 stage 都会将分辨率缩小一倍，channel 数扩大一倍 （like vgg） 。
• Swin Tranformer Block 像大多数 Transformer Block 一样，不改变输入特征 shape，可以看做是一种比较高级（加入了 self-attention）的激活函数。
• 图中 Swin Tranformer Block 都是以连续偶数次出现，因为是一个 W-MSA Transformer Block + 一个 SW-MSA Transformer Block，如右边子图 b 所示。

Patch Partition

• 作用与 ViT 第一步很像：将图片切成若干等大小不重叠的 patch，patch_size = P，然后把每个 patch 从 (P, P, c) 拉成 1 维。
• 本实验中 patch_size = 4，所以一张图被裁切成了 $\frac{H}{4}\times\frac{W}{4}$ 个长度为 4 * 4 * 3 = 48 的向量。

Linear Embedding

• 与一个 shape = (48, C) 的矩阵乘将 48 维映射到 C 维。

与上一步结合可以变成一个 in_channel = 3, out_channel = C, kernel_size = stride = 4 的 Conv2d，官方实现中实际上也是这么做的（ class PatchEmbed ）。

Patch Merging

• 由两步组成，作用是 将分辨率缩小一倍，channel 扩大一倍。
  • 将 $\frac{H}{4}\times\frac{W}{4}\times C$ 的输入使用 bayer2rggb (space2depth with block_size = 2) 变成 $\frac{H}{8}\times\frac{W}{8}\times 4C$ 。
  • 再将 $\frac{H}{8}\times\frac{W}{8}\times 4C$ 与一个 $4C\times 2C$ 的矩阵相乘，输出一个 $\frac{H}{8}\times\frac{W}{8}\times 2C$ 的矩阵。

W-MSA

• 全称为 Windows Multi-head Self-attention，目的是为了减少 Self-attention 计算量。
  • 原始的 MSA 会直接对完整输入计算其 self-attention 结果
  • W-MSA 先将输入拆成大小为 $M\times M$ 且互不重叠的 windows，然后计算每个 window 的 self-attention 结果。
  • 计算公式还是: $Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$
• MSA 和 W-MSA 计算量对比（假设输入特征图 $X \in \mathbb{R}^{H\times W\times C}$，且 $W_q,W_k,W_v \in \mathbb{R}^{C\times C}$，且 Multi-head 的 head 数为 1)：

  • MSA 计算量：

    • X -> Q / K / V 计算量：$\mathbb{R}^{H\times W\times C}\times \mathbb{R}^{C\times C}$ 计算量为 $3HWC^2$ 。

    • $QK^T$ 计算量：$\mathbb{R}^{HW\times C}\times \mathbb{R}^{C\times HW}$ 计算量为 $H^2W^2C$ 。

    • 不考虑 softmax 和 $\frac{..}{\sqrt{d_k}}$ 计算量。

    • softmax 结果 $\times V$ 计算量： $\mathbb{R}^{HW \times HW}\times \mathbb{R}^{HW \times C}$ 计算量为 $H^2W^2C$ 。

    • 因为需要 Multi-head，所以需要将 $\times V$ 之后的矩阵再 $\times W_o$，且 head 数为 1，计算量： $\mathbb{R}^{1*HW \times C}\times \mathbb{R}^{C \times C}$ 计算量为 $HWC^2$ 。

    • 总计算量： $4HWC^2 + 2H^2W^2C$ 。

  • W-MSA 计算量：

    • 上图的 $H\rightarrow M,\ W\rightarrow M$，一共重复 $\frac{HW}{M^2}$ 次，所以总计算量为：

    $(4M^2C^2 + 2M^4C)\times \frac{HW}{M^2}=4HWC^2+2HWM^2C$

  • 相比之下 W-MSA 会比 WSA 计算量少： $2HWC(HW-M^2)$ 。

SW-MSA

• 为了节省计算量 W-MSA 会将输入的完整特征图分 window，每个 window 独立去做 self-attention，这会导致 window 之间的关联性消失，这有悖于 self-attention 会在全图上建立长距离全局相关性依赖的特点。
• 所以，作者引入 SW-MSA 的方案通过 滑窗 解决这一问题。
  
• 滑窗会带来边角处的零碎区域（长或宽小于 window_size 的区域），由于 $H \% M = 0,\ W \% M = 0$ (W-MSA 要求)，所以零碎的区域可以通过调整位置拼成完整 window 。
• 完整的滑窗区域与 W-MSA 一样独立做 Self-attention。
• 由零碎区域拼成的滑窗区域中：

  • 本属于同一零碎区域的位置可以 Self-attention 。

  • 本不属于同一零碎区域的位置在原图上不相邻，不能做 Self-attention 。

  • 具体做法是来自不同区域位置之间的 Self-attention 被 mask 掉 ，只保留来自相同区域的位置间的 Self-attention 。

W-MSA + SW-MSA

• 如网络结构图子图 b 所示，W-MSA 连接 SW-MSA 之后，数学表示：
  $\hat{z}^l=W-MSA(LN(z^{l-1})) + z^{l-1}$
  $z^l=MLP(LN(\hat z^l)) + \hat z^l$
  $\hat z^{l+1} = SW-MSA(LN(z^l))+z^l$
  $z^{l+1} = MLP(LN(\hat z^{l+1}))+\hat z^{l+1}$

Relative Position Bias

• 对 Self-attention 加上 相对位置偏置 bias 之后，精度提升巨大。
• 原始 Self-attention： $Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$ 。
• Relative Position Bias Self-attention： $Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}} + Bias)V$ 。
  

Through

• 质疑了很多人都没有质疑的点，例如给 Self-attention 加上相对位置偏置。
• 本质是对 Tranformer 的一种很 work 的加速方法。
• 但 Shift window 逻辑比较复杂，没有经典模型该有的简约美，盲猜之后会被更 make sense 的结构替代。

Python 闭包 (closure)

闭包定义

• 闭包: 在一些语言中，在函数中可以（嵌套）定义另一个函数时，如果内部的函数引用了外部的函数的变量，则可能产生闭包。闭包可以用来在一个函数与一组“私有”变量之间创建关联关系。在给定函数被多次调用的过程中，这些私有变量能够保持其持久性。
• 支持将函数当成对象使用的编程语言，一般都支持闭包。比如 Python, JavaScript。

闭包的示例

• 代码

  1 2 3 4 5 6 7 8 9 10 11 12 13

  def function_1(arg_1): def function_2(arg_2): return arg_1 * arg_2 return function_2 times_8 = function_1(8) out = times_8(9) print(f"times_8(9) = {out}") # 闭包中的 cell print(f"times_8.__closure__ = {times_8.__closure__}") # 闭包中的 cell 对象的内容 print("times_8.__closure__.cell_contents:") for i in times_8.__closure__: print(i.cell_contents)

• 输出

  1 2 3 4

  times_8(9) = 72 times_8.__closure__ = (<cell at 0x7ff39d4d2a30: int object at 0x5642a56c5e20>,) times_8.__closure__.cell_contents: 8

闭包的用处

1. 可以读取函数内部的变量

• 如上面给出的例子

2. 让这些变量的值始终保持在内存中

• 例如一个棋盘游戏，棋子每次可以选择上下左右方向中的一个，在此方向上移动距离 step，使用闭包实现代码如下：

  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17

  def create(pos=[0, 0]): def go(direction, step): new_x = pos[0] + direction[0] * step new_y = pos[1] + direction[1] * step pos[0] = new_x pos[1] = new_y return pos return go player = create() print(player([1, 0], 10)) print(player([0, 1], 20)) print(player([-1, 0], 10))

• 输出

  1 2 3

  [10, 0] [10, 20] [0, 20]

棋子每次更新后的位置都会存储在闭包中。

3. 用于装饰器

• 可以读取函数内部的变量 和 让这些变量的值始终保持在内存中 都可以使用 Python 的类实现，但 装饰器 是闭包的一个典型用处。

装饰器 (Decorators)

装饰器的定义

• 装饰器: 由闭包的概念引申而来，是一种 增加函数或类功能的方法，它可以快速地给不同的函数或类传入相同的功能。
• 直白点说就是被装饰的函数或类，会作为参数传入到装饰器对应的函数或类中，在装饰器中会对传入的函数或类（被装饰的对象）进行一些处理，然后返回一个新的函数或类（通常是对原函数或类的增强版）。

函数装饰器的示例

• 代码

  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19

  import time def count_time(some_fun): def wrapper(): t1 = time.time() some_fun() print(f"运行时间为: {round(time.time() - t1, 2)} s") return wrapper # 装饰器语法糖 @count_time def function_1(): time.sleep(1) print("run function_1") function_1() # 不使用语法糖的方法 def function_2(): time.sleep(1) print("run function_2") new_function = count_time(function_2) new_function()

• 输出

  1 2 3 4

  run function_1 运行时间为: 1.0 s run function_2 运行时间为: 1.0 s

类别装饰器示例

• 代码

  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21

  import time class Timer: def __init__(self, func) -> None: self.func = func def __call__(self, *args, **kwargs): start = time.time() ret = self.func(*args, **kwargs) print(f"Time: {time.time()-start}") return ret # 使用装饰器语法糖实现 @Timer def add_1(a, b): time.sleep(1) print(f"{a} + {b} = {a+b}") add_1(2, 3) # 不使用装饰器语法糖实现 def add_2(a, b): time.sleep(1) print(f"{a} + {b} = {a+b}") new_add_2 = Timer(add_2) new_add_2(2, 3)

• 输出

  1 2 3 4

  2 + 3 = 5 Time: 1.0011768341064453 2 + 3 = 5 Time: 1.001098394393921

URL

• https://arxiv.org/pdf/2106.08295.pdf

TL;DR

• 本文用比较简洁的方式给出了神经网络的通用量化方法，是量化领域的必读论文。

Algorithm

1. 量化基础知识

1.1 硬件背景

• 一个 $y=Wx+b$ 实际上是由 乘法器 和 累加器 组合而成的，实际的计算过程如下：
  

卷积实际上也是通过 image to column 操作变成 $y=Wx+b$ 操作

• 常见的 int8 量化会将上述过程变成如下过程：
  

weight 和 input 都被量化为 int8 ，同时保留各自的量化 scale，乘法操作是整形乘法器（更快），累加器是 int32 类型，最后再量化为 int8 放到 OCM 上

1.2 均匀仿射量化

• 均匀仿射量化也被称为 非对称量化，由三个量化参数定义：
  • 比例因子 scale
  • 零点 zero_point
  • 比特宽度 bits
• 非对称量化：
  • for unsigned integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil+z;0,2^b-1)$
  • for signed integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil+z;-2^{b-1},2^{b-1}-1)$
  • 这里的 $\lfloor\rceil$ 表示 round 运算
• 对称量化是非对称量化的简化版本，是将零点 zero_point 固定为 0
• 对称量化：
  • for unsigned integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil;0,2^b-1)$
  • for signed integers: $X_{int} = clamp(\lfloor\frac{X}{s}\rceil;-2^{b-1},2^{b-1}-1)$
• 对称量化和非对称量化的含义：
  
  
• 2 的指数幂量化：
  • 限制 $s=2^{-k}$
  • 优势：scale 过程变成了硬件移位，对硬件更友好。
  • 劣势：会使得 round 和 clip 误差的权衡变难。
• 量化颗粒度：
  • per-tensor: 硬件更友好，但限制了量化的自由度。
  • per-channel: 反之。

1.3 量化模拟

• 量化模拟是指在浮点计算设备上模拟定点计算设备的过程，通常用于训练。
  

左边是定点计算过程，右边是用浮点设备模型定点计算的过程

• 为了减少数据搬运和不必要的量化步骤，通常会做：
  • batch norm 折叠：batch norm 在推理时是静态的，因此可以和前面的 conv 等层合并。
  • 激活函数融合：在实际的硬件解决方案中，通常会在非线性操作（如 ReLU）之后直接进行量化，而不是先将激活写入内存然后再加载回计算核心。

1.4 实践考量

• 对称量化和非对称量化：
  • 对称量化：zero-point == 0
  • 非对称量化：zero-point != 0
• 为了方便计算，通常情况下，会将权重设置为对称量化（$z_w=0$），将特征设置为非对称量化（$z_x\ne 0$）
  • 原因分析：
    • $W=S_w(W_{int} - Z_w)$
    • $X=S_x(X_{int} - Z_x)$
    • $WX=S_wS_x(W_{int} - Z_w)(X_{int} - Z_x)\\=S_wS_xW_{int}X_{int}-S_wS_xZ_wX_{int}-S_wS_xZ_xW_{int}+S_wS_xZ_wZ_x$
    • 在推理阶段：$S_w,\ S_x,\ Z_w,\ Z_x,\ W_{int}$ 已知，因此：
      • 等式的第三项和第四项可提前算出，无需推理耗时。
      • 第一项和第二项由于关联动态输入 $X_{int}$，因此需要额外耗时；但是如果设置 $Z_w=0$，则第二项恒等于0，可节省计算量。

2. 训练后量化（PTQ，post-training quantization）

• 训练后量化是指用 float32 精度训练的模型直接转成量化模型，无需任何数据和训练。

2.1 量化范围的设置

• 最大最小值法（min-max）：$q_{min}=minV,\ \ q_{max}=maxV$，$V$ 是待量化 tensor
• 均方差法（MSE）：$\argmin_{q_{min},q_{max}}||V-\hat{V}(q_{min}, q_{max})||^2_F$
• 交叉熵法（cross entropy）：$\argmin_{q_{min},q_{max}}=H(softmax(V),softmax(\hat{V}(q_{min},q_{max})))$，其中 $H$ 表示 cross entropy function
• 批量归一化法（BN based）：$q_{min}=min(\beta-\alpha\gamma),\ \ q_{max}=max(\beta+\alpha\gamma)$，其中 $\beta,\ \gamma$ 分布表示 batch norm 学到的 per channel 的 shift 和 scale，$\alpha>0$ 是超参数
• 组合法（comparsion）：以上方法的自由组合
  
  

使用不同量化方法分别量化 weight 和 activation 后的精度

2.2 跨层均衡（Cross-Layer Equalization）

• 这是一种 通过修改模型权重 来改善神经网络量化性能的技术，CLE 的目的是减少网络中不同 channel 之间由于量化引起的性能不平衡，这种问题在 depth-wise conv layer 中尤其容易出现。
  

mobilenetv2 第一个 depth-wise conv 层的 per output channel weight range

• 想要实现跨层均衡的模型，需要激活函数满足交换律，即：$f(sx)=sf(x)$，常见的 ReLU 和 PReLU 都满足。
  
• CLE 原理：
  • $y=f(W_2(W_1x+b_1)+b_2)\\=f(SW_2(S^{-1}W_1x+S^{-1}b_1)+b_2)\\=f(\hat{W_2}(\hat{W_1}x+\hat{b_1})+b_2)$
  • 其中：
    • $\hat{W_1}=S^{-1}W_1$
    • $\hat{b_1}=S^{-1}b_1$
    • $\hat{W_2}=SW_2$
    • $S_i=\frac{\sqrt {r_i^1r_i^2}}{r_i^2}$，其中 $r_i^j$ 表示 j tensor 的 i channel
• abosrbing high bias 是一种 解决模型中过大 bias 的技术，原理是：
  • $y=W_2h+b_2\\=W_2(f(W_1x+b_1))+b_2\\=W_2(f(W_1x+b_1)+c-c)+b_2\\=W_2(f(W_1x+\hat{b_1})+c)+b_2\\=W_2(f(W_1x+\hat{b_1}))+\hat{b_2}\\=W_2\hat{h}+\hat{b_2}$
  • 其中：
    • $\hat{b_2}=b_2+W_2c$
    • $\hat{h}=h-c$
    • $\hat{b_1}=b_1-c$
    • $c_i=max(0, min_x(W_{1i}x+b_{1i}))$

WIP

URL

https://mlc.ai/zh/chapter_tensor_program/index.html

元张量函数

• 元张量函数 表示机器学习模型计算中的单个单元计算。
  • 一个机器学习编译过程可以有选择地转换元张量函数的实现。

张量程序

• 张量程序 是一个表示元张量函数的有效抽象。
  • 关键成分包括: 多维数组，循环嵌套，计算语句。
  • 程序变换可以被用于加速张量程序的执行。
  • 张量程序中额外的结构能够为程序变换提供更多的信息。

TensorIR: 张量程序抽象案例研究

• TensorIR 是标准机器学习编译框架 Apache TVM 中使用的张量程序抽象。

目标

• 使用 TensorIR 张量程序抽象 ReLU(A @ B) 张量函数。
• 数学表示：
  • $Y_{i,j}=\sum_k A_{i,k}\times B_{k,j}$
  • $C_{i,j}=ReLU(Y_{i,j})=max(Y_{i,j}, 0)$

不同实现方法

使用 Numpy 实现

1
2
3
4
5

dtype = "float32"
a_np = np.random.rand(128, 128).astype(dtype)
b_np = np.random.rand(128, 128).astype(dtype)
# a @ b is equivalent to np.matmul(a, b)
c_mm_relu = np.maximum(a_np @ b_np, 0)

使用 Low Level Numpy 实现

Low Level Numpy 是指只使用 Numpy 的数据结构而不调用 Numpy 的 API

1
2
3
4
5
6
7
8
9
10
11
12

# Use low level numpy to implement matmal ReLU
def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
    Y = np.empty((128, 128), dtype="float32")
    for i in range(128):
        for j in range(128):
            for k in range(128):
                if k == 0:
                    Y[i, j] = 0
                Y[i, j] = Y[i, j] + A[i, k] * B[k, j]
    for i in range(128):
        for j in range(128):
            C[i, j] = max(Y[i, j], 0)

使用 TensorIR 实现

TensorIR 是 TVMScript 中的一种 Python 方言

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
@tvm.script.ir_module
class MyModule:
    @T.prim_func
    def mm_relu(A: T.Buffer[(128, 128), "float32"],
                B: T.Buffer[(128, 128), "float32"],
                C: T.Buffer[(128, 128), "float32"]):
        T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})
        Y = T.alloc_buffer((128, 128), dtype="float32")
        for i, j, k in T.grid(128, 128, 128):
            with T.block("Y"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j)
                vk = T.axis.reduce(128, k)
                with T.init():
                    Y[vi, vj] = T.float32(0)
                Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
        for i, j in T.grid(128, 128):
            with T.block("C"):
                vi = T.axis.spatial(128, i)
                vj = T.axis.spatial(128, j)
                C[vi, vj] = T.max(Y[vi, vj], T.float32(0))

TensorIR 代码与 Low Level Numpy 代码对比

函数参数

1
2
3
4
5
6
7
8

# TensorIR
def mm_relu(A: T.Buffer[(128, 128), "float32"],
            B: T.Buffer[(128, 128), "float32"],
            C: T.Buffer[(128, 128), "float32"]):
    ...
# numpy
def lnumpy_mm_relu(A: np.ndarray, B: np.ndarray, C: np.ndarray):
    ...

buffer

1
2
3
4

# TensorIR
Y = T.alloc_buffer((128, 128), dtype="float32")
# numpy
Y = np.empty((128, 128), dtype="float32")

循环

1
2
3
4
5
6

# TensorIR
for i, j, k in T.grid(128, 128, 128):
# numpy
for i in range(128):
    for j in range(128):
        for k in range(128):

计算块

1
2
3
4
5
6
7
8
9
10
11
12
13

# TensorIR
with T.block("Y"):
    vi = T.axis.spatial(128, i)
    vj = T.axis.spatial(128, j)
    vk = T.axis.reduce(128, k)
    with T.init():
        Y[vi, vj] = T.float32(0)
    Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]
# coressponding numpy code
vi, vj, vk = i, j, k
if vk == 0:
    Y[vi, vj] = 0
Y[vi, vj] = Y[vi, vj] + A[vi, vk] * B[vk, vj]

块（Block） 是 TensorIR 中的基本计算单位。

• 值得注意的是，对于一组固定的 vi 和 vj，计算块在 Y 的空间位置 (Y[vi, vj]) 处生成一个点值，该点值独立于 Y 中的其他位置（具有不同的vi, vj 值的位置）。我们可以称 vi、vj 为 空间轴，因为它们直接对应于块写入的缓冲区空间区域的开始。 涉及归约的轴（vk）被命名为 归约轴。
• 空间轴上的每个点都独立于其他点。

1
2
3
4
5
6

vi = T.axis.spatial(128, i)
vj = T.axis.spatial(128, j)
vk = T.axis.reduce(128, k)
# 使用语法糖可等价写成如下形式
# SSR means the properties of each axes are "spatial", "spatial", "reduce"
vi, vj, vk = T.axis.remap("SSR", [i, j, k])

函数属性

1

T.func_attr({"global_symbol": "mm_relu", "tir.noalias": True})

其中：

• global_symbol 对应函数名。
• tir.noalias 是一个属性，表示所有的缓冲存储器不重叠。

装饰器

• @tvm.script.ir_module 表示被装饰的类是一个 IRModule。
• @T.prim_func 表示被装饰的函数是一个张量函数。

URL

https://mlc.ai/zh/chapter_introduction/index.html

什么是机器学习编译

• 机器学习编译 (machine learning compilation, MLC) 是指，将机器学习算法从开发阶段，通过变换和优化算法，使其变成部署状态。
  

机器学习的痛点之一是：训练框架繁多/部署终端种类繁多，开发与部署存在 gap 。

• 开发形式 是指我们在开发机器学习模型时使用的形式。典型的开发形式包括用 PyTorch、TensorFlow 或 JAX 等通用框架编写的模型描述，以及与之相关的权重。
• 部署形式 是指执行机器学习应用程序所需的形式。它通常涉及机器学习模型的每个步骤的支撑代码、管理资源（例如内存）的控制器，以及与应用程序开发环境的接口（例如用于 android 应用程序的 java API）。
• 机器学习编译的目标
  • 集成与最小化依赖
  • 利用硬件加速
  • 通用优化

机器学习编译的关键要素

• 张量
• 张量函数
• 抽象：做什么
• 实现：怎么做
  

绿色节点表示张量，白色节点表示张量函数

机器学习编译过程中的张量函数变换过程

抽象和实现

总结

• 机器学习编译的目标
  • 集成与最小化依赖
  • 利用硬件加速
  • 通用优化
• 为什么学习机器学习编译
  • 构建机器学习部署解决方案
  • 深入了解现有机器学习框架
  • 为新兴硬件建立软件栈
• 机器学习编译的关键要素
  • 张量和张量函数
  • 抽象和实现是值得思考的工具