Zhangzhe's Blog

Learning Diverse and Discriminative Representations via the Principle of Maximal Coding Rate Reduction

Posted on 2022-12-25 Edited on 2025-03-03 In MCR2 Valine:

URL

paper: https://arxiv.org/pdf/2006.08558.pdf
code: https://github.com/ryanchankh/mcr2.git

Algorithm

TL;DR

本文提出一种 最大编码率降低（MCR2） 表征算法，本质是一种表征损失函数，本算法有效优化了表征空间，在有监督学习（分类）与自监督学习（聚类）都取得了不错的效果。

Maximal Coding Rate Reduction

什么是一个好的表征？一个好的表征应该有哪些性质？
- 一个好的表征应该充分利用表征空间。
- 一个好的表征在同一类下的表征应该尽可能的相似。
$MCR^2$ $M C R^{2}$ Loss
- $R(Z, \epsilon)=\frac{1}{2}logdet(I+\frac{d}{m\epsilon^2}ZZ^T)$
- $R^c(Z,\epsilon|\Pi)=\sum_{j=1}^k\frac{tr(\Pi_j)}{2m}logdet(I+\frac{d}{tr(\Pi_j)\epsilon^2}Z\Pi_j Z^T)$
- $\max_{\theta|\Pi} \Delta R(Z(\theta),\Pi,\epsilon)=R(Z, \epsilon)-R^c(Z,\epsilon|\Pi)$
- 其中：
  - $Z\in\mathbb R^{d\times m}$ ，其中 $d$ 是表征向量的长度， $m$ 是一个 batch 的大小，典型值是 $d=128,\ m=1000$ 。
  - $det$ 表示行列式的值， $logdet$ 表示行列式的值的自然对数。
  - $\Pi_j$ 表示选择函数，选择属于类 $j$ 的特征向量进行计算。
$MCR^2$ $M C R^{2}$ Loss 解析
- $det$ 行列式函数可以用于衡量一个矩阵中向量的正交程度，行列式的值越大，矩阵中向量越正交，向量实际利用的表征空间越大。
- $A=\begin{bmatrix} a & b \\ c & d \end{bmatrix}$ 矩阵的行列式表示由向量 $v_1 = [a, b],\ v_2=[c, d]$ 组成的平行四边形的面积，如下图所示，当 $v_1,\ v_2$ 向量正交时，面积最大，行列式值最大。
- 同理，在 $d$ 维空间下，当 $d$ 个 $d$ 维空间越正交，组成的 空间积 越大。
- $ZZ^T$ 是个实对称矩阵，因此是半正定矩阵，因此 $I + ZZ^T$ 是个正定矩阵，因此 $I+ZZ^T$ 的行列式的值 > 0，因此 logdet 有定义。
- $MCR^2$ Loss 的实际含义是：所有表征向量尽可能正交，属于同一个类的表征向量尽可能不正交，因此属于同一个类别的表征向量会尽可能共线，不同类别会尽可能正交。
- Loss 中的 $d,\ m$ 都是平衡因子，平衡因向量的长度和统计集大小引起的数值变化。
在使用 Cifar10 数据集训练后，将输出的 128 维度表征使用任意分类器（SVM / KNN / 单层神经网络）都很容易进行分类，达到 95+ 的准确率。
而且 $MCR^2$ Loss 在分类任务中的一个优势在于：对于存在错误标签的数据， $MCR^2$ 比交叉熵对错误标签的敏感度更低，如下图所示：
面对聚类任务，由于没有类别信息，损失函数变成： $\max_{\theta|\Pi} \Delta R(Z(\theta),\Pi,\epsilon)=R(Z, \epsilon)$ ，即：尽可能充分利用表征空间 。

Throught

本文提出的方法从表征角度讲非常 make sense，但存在的问题是：依然无法摆脱 维度灾难，因此 $MCR^2$ 也仅仅被用于低维度表征空间中，无法在神经网络的每一层都使用，在分类任务中也仅仅可以被当做一个在交叉熵的升级版本（交叉熵作用于类别维度， $MCR^2$ 监督维度更高）。
一个简单的想法确实可以有效提高聚类任务的模型效果，所以为后面的 Deep (Convolution) Networks from First Principles 提供了理论基础。

CUDA 基础之矩阵乘优化

Posted on 2022-11-21 Edited on 2025-03-03 In CUDA Valine:

CUDA（Compute Unified Device Architecture，统一计算设备架构）资料：

GPU 体系结构

物理模型
- 典型的 GPU 包含一组流处理器 (stream multi-processors, SM)，每个流处理器都有许多核心，硬件实现上这些核心之间可共享内存（shared memory）
逻辑模型
- 逻辑模型中，引入了 Grid / Block / Thread 三级概念，逻辑模型与物理的对应关系如下：
  
  因此：同一个 Block 中的 Thread 可共享 shared memory
Memory Hierarchy

shared memory 速度几乎和 L1 cache 一样，比 local memory 和 global memory 都快的多（在物理上，local memory 和 global memory 是同一块 DRAM）
在对 GPU 进行编程时，需要创建一组进程块 (thread blocks)，每个 thread 映射到单个核心，而 block 映射到流式多处理器 (SM)，如下图所示：
每个线程可由 threadIdx 和 blockIdx 索引，在实际应用中，可以有多维线程索引

共享内存优化

以矩阵乘为例， $A\in \mathbb{R}^{1024\times 1024},B\in \mathbb{R}^{1024\times 1024}$ $A \in R^{1024 \times 1024}, B \in R^{1024 \times 1024}$
- 同一个 block 中的多个 thread 可共享内存，因此可以重排同一个 block 中的 thread 数据，使得尽可能少的数据缓存到 shared memory 中
- 优化前：
  - 每个 thread 需要计算输出矩阵中 8 * 8 的数据，需要从 local memory 中读取 8 * 8 * 1024 * 2 数据
  - 每个 block 中的 thread 之间没有数据共享，所以需要从 local memory 中读取 $8 * 8 * 8 * 8 * 1024 * 2 = 2^{23}$ 个矩阵元素
- 优化后：
  - 每个 block 计算输出矩阵的 64 * 64 的数据最少需要 $64 * 1024 * 2=2^{17}$ 的数据，可提前将这部分数据缓存到 shared memory
  - 然后每个 thread 从 shared memory 读数据计算，需读取 $64 * 1024 * 2=2^{17}$ 个数据
- 内存优化前后每个 block 读取数据对比：
  - 优化前：从 local memory 读取 $2^{23}$ 个矩阵元素
  - 优化后：从 local memory 读取 $2^{17}$ 个矩阵元素到 shared memory，再从 shared memory 读取 $2^{17}$ 个数据计算

机器学习编译总结

Posted on 2022-11-21 Edited on 2025-03-03 In 机器学习编译 Valine:

TVM 是什么：是 Tensor Virtual Machine 的缩写，是一个 Open Deep Learning Compiler Stack（深度学习开源编译栈）
TVM 想干什么：将机器学习算法从开发阶段形态，通过变换和优化算法，使其变成部署形态
TVM 的原则：
- 集成与最小化依赖
- 利用硬件加速
- 通用优化
TVM Module 层次结构：
- IRModule：包含一个或多个 元张量函数 和一个 高层神经网络执行的抽象。通常用 @tvm.script.ir_module 装饰器装饰
- tensorIR: 元张量函数。通常表示一个算子实例的计算过程，包含多个 计算块。通常用 @T.prim_func 装饰器装饰
- 高层神经网络执行的抽象：IRModule 的程序入口。通常用 @R.function
- block: 计算块。张量的基本计算单位，通常包含多个 计算轴 上的循环。通常用 with T.block(block_name) 来标明作用域
- 计算轴：
  - 空间轴（spatial axis）：空间轴上循环的每个位置的计算独立于其他位置
  - 规约轴（reduce axis）：规约轴上的位置不会反映到最后的计算输出上
TVM Module 变换过程：
- 自动程序优化
- cuda 多线程优化
- 内存优化
- 图优化
- 等等

TVM Module 执行过程：

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)
虚拟机执行器 = 虚拟机(可执行程序)
运行结果 = 虚拟机执行器(模型输入 + 模型权重)

机器学习编译(6)——GPU硬件加速

Posted on 2022-11-09 Edited on 2025-03-03 In 机器学习编译 Valine:

URL

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

背景知识

CUDA（Compute Unified Device Architecture，同一计算设备架构）教程：

GPU 体系结构

物理模型
- 典型的 GPU 包含一组流处理器 (stream multi-processors, SM)，每个流处理器都有许多核心，硬件实现上这些核心之间可共享内存（shared memory）
逻辑模型
- 逻辑模型中，引入了 Grid / Block / Thread 三级概念，逻辑模型与物理的对应关系如下图所示：
  
  因此：同一个 Block 中的 Thread 可共享 shared memory
Memory Hierarchy

shared memory 速度几乎和 L1 cache 一样，比 local memory 和 global memory 都快的多（在物理上，local memory 和 global memory 是同一块 DRAM）
在对 GPU 进行编程时，需要创建一组进程块 (thread blocks)，每个 thread 映射到单个核心，而 block 映射到流式多处理器 (SM)，如下图所示：
每个线程可由 threadIdx 和 blockIdx 索引，在实际应用中，可以有多维线程索引

Element-wise Add GPU 加速

两个向量 A 和 B，向量长度都为 1024,执行元素相加，并将结果存储在 C 中

@tvm.script.ir_module
class MyModuleVecAdd:
    @T.prim_func
    def main(A: T.Buffer[(1024,), "float32"],
            B: T.Buffer[(1024,), "float32"],
            C: T.Buffer[(1024,), "float32"]) -> None:
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        for i in T.grid(1024):
            with T.block("C"):
                vi = T.axis.remap("S", [i])
                C[vi] = A[vi] + B[vi]

首先将循环 i 拆分成两个循环:

sch = tvm.tir.Schedule(MyModuleVecAdd)
block_C = sch.get_block("C")
i, = sch.get_loops(block=block_C)
i0, i1 = sch.split(i, [None, 128])

将迭代器绑定到 GPU 线程块。每个线程由两个索引进行表示 threadIdx.x 和 blockIdx.x
1
2
3
sch.bind(i0, "blockIdx.x")
sch.bind(i1, "threadIdx.x")
sch.mod.show()

绑定后的代码：

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[1024, "float32"], B: T.Buffer[1024, "float32"], C: T.Buffer[1024, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i_0 in T.thread_binding(8, thread="blockIdx.x"):
            for i_1 in T.thread_binding(128, thread="threadIdx.x"):
                with T.block("C"):
                    vi = T.axis.spatial(1024, i_0 * 128 + i_1)
                    T.reads(A[vi], B[vi])
                    T.writes(C[vi])
                    C[vi] = A[vi] + B[vi]

由于 Element-wise Add 不存在数据依赖，所以可以直接拆分到多个 block 中的多个 thread 中，一个 cycle 全部算完

窗口求和 GPU 加速

相邻三个窗口求和，输入向量 A 长度 1026，输出 B 长度 1024。（即无 padding 的权重为 [1, 1, 1] 的 conv1d）

@tvm.script.ir_module
class MyModuleWindowSum:
    @T.prim_func
    def main(A: T.Buffer[(1026,), "float32"],
            B: T.Buffer[(1024,), "float32"]) -> None:
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        for i in T.grid(1024):
            with T.block("C"):
                vi = T.axis.remap("S", [i])
                B[vi] = A[vi] + A[vi + 1] + A[vi + 2]

拆分循环并绑定到 block 和 thread

sch = tvm.tir.Schedule(MyModuleWindowSum)
nthread = 128
block_C = sch.get_block("C")
i,  = sch.get_loops(block=block_C)
i0, i1 = sch.split(i, [None, nthread])
sch.bind(i0, "blockIdx.x")
sch.bind(i1, "threadIdx.x")
sch.mod.show()

拆分循环后 IRModule

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[1027, "float32"], B: T.Buffer[1024, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        for i_0 in T.thread_binding(8, thread="blockIdx.x"):
            for i_1 in T.thread_binding(128, thread="threadIdx.x"):
                # 启用 8 个 block 并发计算，每个 block 用 16 个 thread 并发
                # 因此每一个 thread 只需要计算 1 次乘加
                with T.block("C"):
                    vi = T.axis.spatial(1024, i_0 * 128 + i_1)
                    T.reads(A[vi : vi + 3])
                    T.writes(B[vi])
                    B[vi] = A[vi] + A[vi + 1] + A[vi + 2]

提前缓存数据

1
2
3

A_shared = sch.cache_read(block_C, read_buffer_index=0, storage_scope="shared")
sch.compute_at(A_shared, i1)
sch.mod.show()

提前缓存数据后的 IRModule

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[1027, "float32"], B: T.Buffer[1024, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        A_shared = T.alloc_buffer([1027], dtype="float32", scope="shared")
        for i_0 in T.thread_binding(8, thread="blockIdx.x"):
            for i_1 in T.thread_binding(128, thread="threadIdx.x"):
                # 由上图 GPU 结构图可知
                # 不同 block 无法共享 share memory
                # 相同 block 的不同 thread 之间可以共享
                # 所以输出 128 个结果需要 130 个输入（本行 128 个加下一行 2 个）
                for ax0 in T.serial(130):
                    with T.block("A_shared"):
                        v0 = T.axis.spatial(1027, i_0 * 128 + ax0)
                        T.reads(A[v0])
                        T.writes(A_shared[v0])
                        A_shared[v0] = A[v0]
                with T.block("C"):
                    vi = T.axis.spatial(1024, i_0 * 128 + i_1)
                    T.reads(A_shared[vi : vi + 3])
                    T.writes(B[vi])
                    B[vi] = A_shared[vi] + A_shared[vi + 1] + A_shared[vi + 2]

缓存数据可以使用多线程优化
- 因为内存是跨线程共享的，所以需要重新拆分循环并将获取过程的内部迭代器绑定到线程索引上，这种技术称为 cooperative fetching
1
2
3
4
ax = sch.get_loops(A_shared)[-1]
ax0, ax1 = sch.split(ax, [None, nthread])
sch.bind(ax1, "threadIdx.x")
sch.mod.show()

缓存数据优化后 IRModule

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[1026, "float32"], B: T.Buffer[1024, "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        A_shared = T.alloc_buffer([1026], dtype="float32", scope="shared")
        for i_0 in T.thread_binding(8, thread="blockIdx.x"):
            for i_1 in T.thread_binding(128, thread="threadIdx.x"):
                for ax0_0 in T.serial(2):
                    for ax0_1 in T.thread_binding(128, thread="threadIdx.x"):
                        with T.block("A_shared"):
                            # 由上图 GPU 结构图可知
                            # 不同 block 无法共享 share memory
                            # 相同 block 的不同 thread 之间可以共享
                            # 所以输出 128 个结果需要 130 个输入（本行 128 个加下一行 2 个）
                            v0 = T.axis.spatial(1026, i_0 * 128 + (ax0_0 * 128 + ax0_1))
                            T.where(ax0_0 * 128 + ax0_1 < 130)
                            T.reads(A[v0])
                            T.writes(A_shared[v0])
                            A_shared[v0] = A[v0]
                with T.block("C"):
                    vi = T.axis.spatial(1024, i_0 * 128 + i_1)
                    T.reads(A_shared[vi : vi + 3])
                    T.writes(B[vi])
                    B[vi] = A_shared[vi] + A_shared[vi + 1] + A_shared[vi + 2]

矩阵乘法 GPU 加速

IRModule 基础实现：

@tvm.script.ir_module
class MyModuleMatmul:
    @T.prim_func
    def main(A: T.Buffer[(1024, 1024), "float32"],
            B: T.Buffer[(1024, 1024), "float32"],
            C: T.Buffer[(1024, 1024), "float32"]) -> None:
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        for i, j, k in T.grid(1024, 1024, 1024):
            with T.block("C"):
                vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                with T.init():
                    C[vi, vj] = 0.0
                C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]

绑定 block 和 thread + 本地存储分块优化

循环拆分，来增加整体内存复用，只需要从 A 和 B 加载一次条形数据（上图中的灰色部分），然后使用它们来计算矩阵乘法结果
下面代码中设置 V = 8

def blocking(sch,
            tile_local_y,
            tile_local_x,
            tile_block_y,
            tile_block_x,
            tile_k):
    block_C = sch.get_block("C")
    C_local = sch.cache_write(block_C, 0, "local")
    i, j, k = sch.get_loops(block=block_C)
    i0, i1, i2 = sch.split(loop=i, factors=[None, tile_block_y, tile_local_y])
    j0, j1, j2 = sch.split(loop=j, factors=[None, tile_block_x, tile_local_x])
    k0, k1 = sch.split(loop=k, factors=[None, tile_k])
    sch.unroll(k1)
    sch.reorder(i0, j0, i1, j1, k0, k1, i2, j2)
    sch.reverse_compute_at(C_local, j1)
    sch.bind(i0, "blockIdx.y")
    sch.bind(j0, "blockIdx.x")
    sch.bind(i1, "threadIdx.y")
    sch.bind(j1, "threadIdx.x")
    sch.decompose_reduction(block_C, k0)
    return sch
sch = tvm.tir.Schedule(MyModuleMatmul)
sch = blocking(sch, 8, 8, 8, 8, 4)
sch.mod.show()

输出优化后的 IRModule

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[(1024, 1024), "float32"], B: T.Buffer[(1024, 1024), "float32"], C: T.Buffer[(1024, 1024), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        C_local = T.alloc_buffer([1024, 1024], dtype="float32", scope="local")
        for i_0 in T.thread_binding(16, thread="blockIdx.y"):
            for j_0 in T.thread_binding(16, thread="blockIdx.x"):
                for i_1 in T.thread_binding(8, thread="threadIdx.y"):
                    for j_1 in T.thread_binding(8, thread="threadIdx.x"):
                        # 一共使用 16 * 16 个 block 并发计算
                        # 每个 block 使用 8 * 8 个 thread 并发
                        # 所以每个 thread 只需计算输出为 8 * 8 的区域，因此只需要加载 A 中 8 行和 B 中 8 列数据
                        # 1. 初始化 8 * 8 的输出区域为 0
                        for i_2_init, j_2_init in T.grid(8, 8):
                            with T.block("C_init"):
                                vi = T.axis.spatial(1024, i_0 * 64 + i_1 * 8 + i_2_init)
                                vj = T.axis.spatial(1024, j_0 * 64 + j_1 * 8 + j_2_init)
                                T.reads()
                                T.writes(C_local[vi, vj])
                                C_local[vi, vj] = T.float32(0)
                        
                        # 2. 计算 8 * 8 输出区域的值，共计算 8 * 8 * 1024 次乘加
                        for k_0 in T.serial(256):
                            for k_1 in T.unroll(4):
                                for i_2, j_2 in T.grid(8, 8):
                                    with T.block("C_update"):
                                        vi = T.axis.spatial(1024, i_0 * 64 + i_1 * 8 + i_2)
                                        vj = T.axis.spatial(1024, j_0 * 64 + j_1 * 8 + j_2)
                                        vk = T.axis.reduce(1024, k_0 * 4 + k_1)
                                        T.reads(C_local[vi, vj], A[vi, vk], B[vk, vj])
                                        T.writes(C_local[vi, vj])
                                        C_local[vi, vj] = C_local[vi, vj] + A[vi, vk] * B[vk, vj]
                        
                        # 3. 把每个 thread 的 8 * 8 的输出区域拼成最后的 1024 * 1024 的输出
                        for ax0, ax1 in T.grid(8, 8):
                            with T.block("C_local"):
                                v0 = T.axis.spatial(1024, i_0 * 64 + i_1 * 8 + ax0)
                                v1 = T.axis.spatial(1024, j_0 * 64 + j_1 * 8 + ax1)
                                T.reads(C_local[v0, v1])
                                T.writes(C[v0, v1])
                                C[v0, v1] = C_local[v0, v1]

共享内存优化

与上图不同，图中矩阵 C 中 L * L 灰色区域表示一个 block 的计算输出
每个 L * L 灰色区域由多个 V * V 的小区域组成，表示一个 thread 的输出

同一个 block 中的多个 thread 可共享内存，因此可以重排同一个 block 中的 thread 数据，使得尽可能少的数据缓存到 shared memory 中
优化前：
- 每个 thread 需要计算输出矩阵中 8 * 8 的数据，需要从 local memory 中读取 8 * 8 * 1024 * 2 数据
- 每个 block 中的 thread 之间没有数据共享，所以需要从 local memory 中读取 $8 * 8 * 8 * 8 * 1024 * 2 = 2^{23}$ 个矩阵元素
优化后：
- 每个 block 计算输出矩阵的 64 * 64 的数据最少需要 $64 * 1024 * 2=2^{17}$ 的数据，可提前将这部分数据缓存到 shared memory
- 然后每个 thread 从 shared memory 读数据计算，需读取 $64 * 1024 * 2=2^{17}$ 个数据
内存优化前后每个 block 读取数据对比：
- 优化前：从 local memory 读取 $2^{23}$ 个矩阵元素
- 优化后：从 local memory 读取 $2^{17}$ 个矩阵元素到 shared memory，再从 shared memory 读取 $2^{17}$ 个数据计算
优化过程：

def cache_read_and_coop_fetch(sch, block, nthread, read_idx, read_loc):
    read_cache = sch.cache_read(block=block, read_buffer_index=read_idx, storage_scope="shared")
    sch.compute_at(block=read_cache, loop=read_loc)
    # vectorized cooperative fetch
    inner0, inner1 = sch.get_loops(block=read_cache)[-2:]
    inner = sch.fuse(inner0, inner1)
    _, tx, vec = sch.split(loop=inner, factors=[None, nthread, 4])
    sch.vectorize(vec)
    sch.bind(tx, "threadIdx.x")
def blocking_with_shared(
    sch,
    tile_local_y,
    tile_local_x,
    tile_block_y,
    tile_block_x,
    tile_k):
    block_C = sch.get_block("C")
    C_local = sch.cache_write(block_C, 0, "local")
    i, j, k = sch.get_loops(block=block_C)
    i0, i1, i2 = sch.split(loop=i, factors=[None, tile_block_y, tile_local_y])
    j0, j1, j2 = sch.split(loop=j, factors=[None, tile_block_x, tile_local_x])
    k0, k1 = sch.split(loop=k, factors=[None, tile_k])
    sch.reorder(i0, j0, i1, j1, k0, k1, i2, j2)
    sch.reverse_compute_at(C_local, j1)
    sch.bind(i0, "blockIdx.y")
    sch.bind(j0, "blockIdx.x")
    tx = sch.fuse(i1, j1)
    sch.bind(tx, "threadIdx.x")
    nthread = tile_block_y * tile_block_x
    cache_read_and_coop_fetch(sch, block_C, nthread, 0, k0)
    cache_read_and_coop_fetch(sch, block_C, nthread, 1, k0)
    sch.decompose_reduction(block_C, k0)
    return sch
sch = tvm.tir.Schedule(MyModuleMatmul)
sch = blocking_with_shared(sch, 8, 8, 8, 8, 8)
sch.mod.show()

优化后 IRModule

@tvm.script.ir_module
class Module:
    @T.prim_func
    def main(A: T.Buffer[(1024, 1024), "float32"], B: T.Buffer[(1024, 1024), "float32"], C: T.Buffer[(1024, 1024), "float32"]) -> None:
        # function attr dict
        T.func_attr({"global_symbol": "main", "tir.noalias": True})
        # body
        # with T.block("root")
        C_local = T.alloc_buffer([1024, 1024], dtype="float32", scope="local")
        A_shared = T.alloc_buffer([1024, 1024], dtype="float32", scope="shared")
        B_shared = T.alloc_buffer([1024, 1024], dtype="float32", scope="shared")
        for i_0 in T.thread_binding(16, thread="blockIdx.y"):
            for j_0 in T.thread_binding(16, thread="blockIdx.x"):
                for i_1_j_1_fused in T.thread_binding(64, thread="threadIdx.x"):
                    for i_2_init, j_2_init in T.grid(8, 8):
                        with T.block("C_init"):
                            vi = T.axis.spatial(1024, i_0 * 64 + i_1_j_1_fused // 8 * 8 + i_2_init)
                            vj = T.axis.spatial(1024, j_0 * 64 + i_1_j_1_fused % 8 * 8 + j_2_init)
                            T.reads()
                            T.writes(C_local[vi, vj])
                            C_local[vi, vj] = T.float32(0)
                    for k_0 in T.serial(128):
                        for ax0_ax1_fused_0 in T.serial(2):
                            for ax0_ax1_fused_1 in T.thread_binding(64, thread="threadIdx.x"):
                                for ax0_ax1_fused_2 in T.vectorized(4):
                                    with T.block("A_shared"):
                                        v0 = T.axis.spatial(1024, i_0 * 64 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 8)
                                        v1 = T.axis.spatial(1024, k_0 * 8 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 8)
                                        T.reads(A[v0, v1])
                                        T.writes(A_shared[v0, v1])
                                        A_shared[v0, v1] = A[v0, v1]
                        for ax0_ax1_fused_0 in T.serial(2):
                            for ax0_ax1_fused_1 in T.thread_binding(64, thread="threadIdx.x"):
                                for ax0_ax1_fused_2 in T.vectorized(4):
                                    with T.block("B_shared"):
                                        v0 = T.axis.spatial(1024, k_0 * 8 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) // 64)
                                        v1 = T.axis.spatial(1024, j_0 * 64 + (ax0_ax1_fused_0 * 256 + ax0_ax1_fused_1 * 4 + ax0_ax1_fused_2) % 64)
                                        T.reads(B[v0, v1])
                                        T.writes(B_shared[v0, v1])
                                        B_shared[v0, v1] = B[v0, v1]
                        for k_1, i_2, j_2 in T.grid(8, 8, 8):
                            with T.block("C_update"):
                                vi = T.axis.spatial(1024, i_0 * 64 + i_1_j_1_fused // 8 * 8 + i_2)
                                vj = T.axis.spatial(1024, j_0 * 64 + i_1_j_1_fused % 8 * 8 + j_2)
                                vk = T.axis.reduce(1024, k_0 * 8 + k_1)
                                T.reads(C_local[vi, vj], A_shared[vi, vk], B_shared[vk, vj])
                                T.writes(C_local[vi, vj])
                                C_local[vi, vj] = C_local[vi, vj] + A_shared[vi, vk] * B_shared[vk, vj]
                    for ax0, ax1 in T.grid(8, 8):
                        with T.block("C_local"):
                            v0 = T.axis.spatial(1024, i_0 * 64 + i_1_j_1_fused // 8 * 8 + ax0)
                            v1 = T.axis.spatial(1024, j_0 * 64 + i_1_j_1_fused % 8 * 8 + ax1)
                            T.reads(C_local[v0, v1])
                            T.writes(C[v0, v1])
                            C[v0, v1] = C_local[v0, v1]

程序自动变换

from tvm import meta_schedule as ms
sch_tuned = ms.tune_tir(
    mod=MyModuleMatmul,
    target="nvidia/tesla-p100",
    config=ms.TuneConfig(
      max_trials_global=64,
      num_trials_per_iter=64,
    ),
    work_dir="./tune_tmp",
    task_name="main"
)
sch_tuned.mod.show()

机器学习编译(5)——与机器学习框架的整合

Posted on 2022-10-30 Edited on 2025-03-03 In 机器学习编译 Valine:

URL

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

使用 `Builder` 创建 `IRModule`

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

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()

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

# 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`

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

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）

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` 模型

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)

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

# 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`

# 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

@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` 更高层的算子

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()

输出

@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，实现与现有的机器学习框架的整合

机器学习编译(4)——自动程序优化

Posted on 2022-10-22 Edited on 2025-03-03 In 机器学习编译 Valine:

URL

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

自动程序优化的原因

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

自动程序优化过程

end-to-end 构建模型

# 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 主张量函数

# 调优 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)
# 重新测速，速度变快

机器学习编译(3)——端到端模型执行

Posted on 2022-10-21 Edited on 2025-03-03 In 机器学习编译 Valine:

URL

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

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

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

构造 `IR_Module`

@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) 分配输出内存并和输入一起输入主张量函数

构建并运行模型

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)
虚拟机执行器 = 虚拟机(可执行文件)
运行结果 = 虚拟机执行器(模型输入)

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

@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 函数。

相机内参与外参

Posted on 2022-10-16 Edited on 2025-03-03 In 3D Valine:

三种成像相关坐标系

像素坐标系：一种 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$

Lift, Splat, Shoot: Encoding Images from Arbitrary Camera Rigs by Implicitly Unprojecting to 3D

Posted on 2022-09-30 Edited on 2025-03-03 In BEV Valine:

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 轴正方向为竖直向上。

算法细节

特别细节的看代码

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 背景才能看懂

Swin Transformer: Hierarchical Vision Transformer using Shifted Windows

Posted on 2022-09-29 Edited on 2025-03-03 In Transformer Valine:

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}$ $X \in R^{H \times W \times C}$ ，且 $W_q,W_k,W_v \in \mathbb{R}^{C\times C}$ $W_{q}, W_{k}, W_{v} \in 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 的结构替代。

URL

Algorithm

TL;DR

Maximal Coding Rate Reduction

Throught

CUDA（Compute Unified Device Architecture，统一计算设备架构）资料：

GPU 体系结构

共享内存优化

URL

背景知识

GPU 体系结构

Element-wise Add GPU 加速

窗口求和 GPU 加速

矩阵乘法 GPU 加速

程序自动变换

URL

使用 Builder 创建 IRModule

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

使用 BlockBuilder 构造 IRModule

从 Pytorch 映射到 IRModule

Pytorch 模型

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

映射 Pytorch Module 到 TensorIR

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

总结

URL

自动程序优化的原因

自动程序优化过程

end-to-end 构建模型

自动优化 linear0 主张量函数

URL

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

构造 IR_Module

构建并运行模型

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

绑定参数到 IR_Module

总结

三种成像相关坐标系

相机内参

相机外参

总结

URL

TL;DR

Algorithm

Dataset

相关知识

算法细节

1. 特征提取（Lift）

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

3. 体素编码降维（Shoot）

4. 训练 loss

Thought

URL

TL;DR

Algorithm

总体结构

Patch Partition

Linear Embedding

Patch Merging

W-MSA

SW-MSA

W-MSA + SW-MSA

Relative Position Bias

Through

使用 `Builder` 创建 `IRModule`

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

使用 `BlockBuilder` 构造 `IRModule`

从 `Pytorch` 映射到 `IRModule`

`Pytorch` 模型

映射 `Pytorch Module` 到 `TensorIR`

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

构造 `IR_Module`