AutoTVM/Ansor：强化学习调优

可视化演示

展示强化学习调优机制：状态-动作-环境-奖励循环、策略动作搜索与最优动作选择、RL调优相对基线的训练效果，通过交互理解AutoTVM/Ansor的RL调度搜索原理。

调优困境：组合爆炸的厨房

学徒厨师第一天上班就傻了眼。主厨给的那道菜，光可调的旋钮就有七八个：块的大小、循环的顺序、向量化的宽度、数据摆放的方式……每个旋钮还有十几个档位。他随手算了一下：几千万种组合。就算十秒钟试一次，也要连续试好几年。这还没算上每次试菜都要重新切菜、热锅、摆盘——在真实的调优任务里，每试一种调度方案，都要重新编译一次算子再跑一遍，单次试验往往要花几秒甚至几分钟。

这道菜的困境，本质是“组合爆炸”：每个旋钮的档位数相乘，会得到一个天文数字。传统做法靠资深工程师写启发式规则，比如“在 GPU 上优先用某某分块大小”。但规则是针对某个硬件手工调出来的，换一块新芯片、换一个输入尺寸，规则立刻失效。而且规则本身也不覆盖所有算子结构，很多形状复杂的算子根本套不上现成模板。

学徒垂头丧气地放下菜刀，意识到自己缺的不是力气，而是一套“边做边学”的方法。他决定从明天起，每次做菜都认真记笔记，看看能不能从几千次尝试里总结出规律——这就是强化学习登场的原因。

强化学习框架：试菜记笔记的循环

从第二天开始，学徒换了一种干活方式。他每做一次菜，就掏出一个小本子：第一行写“我看到了什么”——食材种类、切菜台的状态、灶台的型号；第二行写“我做了什么”——选了哪种切法、哪个火候；第三行写“客人打了几分”——这道菜上桌后比正常慢了还是快了。试的次数多了，他慢慢发现：某些动作组合总是拿高分，于是他在本子上把这些动作偷偷圈起来，下次多做。

这就是强化学习调优的完整骨架。在 AutoTVM/Ansor 的场景里，学徒是“策略网络”，厨房是“编译器 + 真实硬件”，本子上的“看到什么”叫状态（state），“做了什么”叫动作（action），“客人打分”叫奖励（reward）。一次完整的做菜过程叫一条轨迹，包含多步“状态→动作→奖励”的交替。调度搜索因此被建模成一个序列决策问题：代理一步一步地选择调度变换，最终拼出一个完整的 schedule。

收集了一堆轨迹之后，怎么用？核心思想是：凡是带来高奖励的动作，就提高它下次被选中的概率；凡是低奖励的动作，就降低概率。这个“提优降劣”的数学操作，叫策略梯度。

学徒的本子越记越厚，动作也从盲目乱试变成了“优先尝试划圈的做法”。接下来他要解决三个更细的问题：如何把厨房的状态记清楚？如何保证动作合法？客人打分怎么打才公平？下面四节逐一拆解。

状态空间设计：给后厨拍一张结构化快照

学徒发现，光记“食材很多、锅很热”这种模糊描述没用，他需要一张标准化的表：菜品是什么类型、食材有多少块、每一块大概多大、灶台是明火还是电磁炉、锅的容量多少。换句话说，不能把厨房的乱象全拍下来，而要抽出影响做菜结果的关键信息。

在算子调优中，状态就是对“当前这个算子和硬件”的结构化描述。它要编码算子的计算图结构：算子类型（卷积、矩阵乘、池化……）、输入张量的形状、循环嵌套的层数与层次关系、当前已选择的部分调度。同时还要融入硬件特征：目标芯片是 CPU 还是 GPU 或 NPU、缓存多大、向量寄存器多宽、核数多少。这些特征共同决定哪个调度方案跑得快。

但原始特征有个麻烦：有的特征数值很大（比如张量长度上万），有的很小（比如向量宽度只有几十）。直接把它们塞给神经网络，大数值会淹没小数值，训练很不稳定。所以要先做归一化，把每个维度的特征压到相近的尺度上，再通过嵌入层把离散结构（比如算子类型）转成稠密向量。

有了这张“结构化快照”，学徒每次进后厨都能一眼看清关键信息。但他还缺一样东西：一份明确的“可以做什么”的清单，不然他可能会把菜扔到天花板上。

动作空间设计：合法做菜手法清单

后厨贴着一张规定：学徒只能使用五种手法——切块、焯水、爆炒、摆盘、换锅，而且使用顺序不能颠倒到离谱。比如不能先摆盘再切块。这张清单既保护了菜品的可行性，也大幅缩小了学徒需要尝试的范围。

对应到调优，动作就是“调度变换原语”。AutoTVM/Ansor 里常见的原语有：split（把循环分块）、reorder（交换循环顺序）、vectorize（向量化）、parallel（并行化）、compute_at（把某段计算挪到另一层循环内）等等。一个完整的调度方案，就是这些原语按某种顺序拼成的动作序列。代理的策略，本质上是在为“下一步该用哪个原语、参数选多少”打分。

但并不是所有动作拼接都合法。例如对已经展开的循环再做 split 可能破坏语义，某些变换顺序在不同算子结构下会生成错误代码。假如代理乱试这些非法动作，轻则编译失败浪费试验次数，重则得到错误结果。解决办法是“合法动作掩码”：每一步先把当前状态允许的动作列出来，把非法动作的概率直接置零，再在合法集合上做 softmax。

学徒终于拥有了“合法手法清单”。但他很快发现一个新问题：客人打分必须公平，否则一次手抖就会误导他好几天。

奖励函数设计：客人的评分标准

学徒一开始让客人随便打分，结果有的客人心情好给满分，有的客人挑剔给零分，他完全学不到规律。后来主厨定了规矩：打分只看一道硬指标——上菜时间。先用最普通的做法做一版，记下耗时作为基准；之后每道菜的分数，就是基准耗时除以这道菜的实际耗时。上菜比基准快一倍，分数就是 2；比基准慢一半，分数就是 0.5。

在调优任务里，这个“客人评分”就是奖励函数。首先用一张默认 schedule 编译并运行算子，测出基线耗时 \(T_{\text{baseline}}\)；然后对候选 schedule 同样编译运行，测出 \(T_{\text{candidate}}\)。奖励取两者的比值：

不过比值奖励也有隐患：运行时间受系统噪声影响，同一个方案测两次可能相差 10%。更麻烦的是“奖励稀疏”问题——在组合空间里，绝大多数随机方案性能都差不多，奖励几乎没有区分度，代理像个没有方向的学徒，试了很多次也学不到东西。工程上常用 log 加速比压缩极端值、对同一方案多次测量取中位数、或者先用粗粒度方法筛掉明显差劲的区域，再让 RL 精调。这样才能让客人的评分既稳定又有信息量。

学习探索策略：守成与尝新

学徒现在面临一个经典的矛盾：已知“爆炒拿高分”，他应该每道菜都爆炒吗？如果永远只做拿手菜，他就永远发现不了可能更好的“低温慢煮”。但若是每道菜都乱试新花样，客人又会纷纷给差评。成熟的学徒会定下规矩：大部分时候按已知好做法来，留一小部分次数尝试新配方，并且每次试新风波不能太大，免得把学过的经验全搞砸。

这个矛盾在强化学习里叫“探索与利用的权衡”。对应策略包括：\(\epsilon\)-greedy（以 \(\epsilon\) 概率随机探索，否则选当前最优）、在损失函数里加熵正则项鼓励多样化动作、以及把采集到的轨迹放进经验池打乱顺序再学习（经验回放，减少相邻轨迹的关联性）。AutoTVM/Ansor 这类框架接入 RL 时，常用 PPO 优化器来更新策略，因为它有一个特别务实的机制：限制单次更新的幅度。

学徒带着“大部分守成、小部分尝新、每次小步改进”的纪律，手艺稳步上涨。现在该跟其他后厨的效率比一比了。

效果对比分析：学徒对决掷骰子

餐厅老板想验证学徒的方法值不值得培养，于是请来两组对照组：一组是“掷骰子厨师”，每次做菜完全随机选配方，从不记笔记；另一组是“进化厨师”，同时养着几十个配方，每轮留好的、杂交差的，像自然选择一样迭代。老板给三组同样的预算——各做 200 次试验，看谁先找到满足加速目标的最佳配方。

结果通常很明显：随机搜索完全没有记忆，这次试了高分区域，下次照样可能在低分区域打转，达成目标需要的试验次数非常多。进化算法靠种群并行，能找到不错的解，但每个“个体”都从零开始，缺乏对“一个算子的调度长什么样”的全局理解。RL 代理则把历史试验全部压缩进策略参数里，随着试验次数增加，采样越来越集中在高奖励区域。量化这个差距的指标叫“样本利用率”：

当然，样本效率高不等于永远赢。当搜索空间很小、几十次试验就能穷举完时，随机搜索连训练成本都不用付，反而更干脆；进化算法天然适合并行，多台机器同时跑能摊薄时间。学徒的胜利是有条件的：空间足够大、每次试验成本足够高、并且有足够的历史经验可以积累时，他的“记笔记式学习”才真正划算。

工程选型建议：何时启用学徒制

餐厅老板最后算了一笔账：培养一个学徒有固定成本——要花时间教他记笔记、训练他的判断力。如果餐厅只有一道凉菜，随便让切菜工试几下就行，没必要上学徒制；但如果要应对几十道硬菜、还要长期经营同一间厨房，学徒制就非常值。

调优工作的工程选型也是同样逻辑。先看三个因素：待调算子的数量 \(n_{\text{op}}\)、每个算子需要的平均试验次数 \(N_{\text{trial}}\)、能并行跑试验的设备数 \(p_{\text{device}}\)。三者决定总调优成本：

实际的工程选型还有几条经验法则：搜索空间小（比如旋钮少于四五个）直接用随机或网格搜索；空间中等用进化算法或模拟退火；空间大且硬件会长期服役，才投入 RL。冷启动阶段可以先加载一个预训练策略（从相似算子迁移经验）再微调，避开几百次“从头学起”的浪费。AutoTVM/Ansor 本身也提供了多种搜索器接口，RL 代理并非要替代一切，而是作为可插拔的“精调专家”，与启发式粗筛、进化并行形成混合策略。学徒制的精髓不在于每道菜都从零教，而在于把经验沉淀下来，让下一道菜学得更快。

AI 应用场景

Python 演示

import numpy as np class TuningEnv:    """模拟一个算子调优环境: 5 个调度旋钮, 存在隐藏最优组合"""    def __init__(self, seed=0):        rng = np.random.default_rng(seed)        self.knob_sizes = [4, 4, 8, 3, 8]        self.optimal = [int(rng.integers(n)) for n in self.knob_sizes]        self.baseline_ms = 10.0     def reward(self, schedule):        # 离最优越远, 耗时越长, 加速比越小        dist = sum(s != o for s, o in zip(schedule, self.optimal))        time_ms = self.baseline_ms * (0.7 + 0.3 * dist)        return self.baseline_ms / max(time_ms, 1e-3)  class PolicyAgent:    """用策略梯度(REINFORCE)学习每个旋钮各档位的选择概率"""    def __init__(self, knob_sizes, lr=0.1):        self.theta = [np.zeros(int(n)) for n in knob_sizes]        self.lr = lr        self.reward_mean = 1.0     def softmax(self, logits):        e = np.exp(logits - logits.max())        return e / e.sum()     def sample(self):        # 按当前策略选一个完整调度, 同时计算对数概率        schedule = []        log_prob = 0.0        for j, logits in enumerate(self.theta):            p = self.softmax(logits)            a = int(np.random.choice(len(p), p=p))            schedule.append(a)            log_prob += np.log(p[a] + 1e-12)        return schedule, log_prob     def update(self, schedule, log_prob, reward):        # 优势 = 奖励 - 平均奖励, 降低方差        advantage = reward - self.reward_mean        self.reward_mean = 0.9 * self.reward_mean + 0.1 * reward        for j, a in enumerate(schedule):            p = self.softmax(self.theta[j])            grad = np.zeros_like(p)            grad[a] += 1.0            grad -= p            self.theta[j] += self.lr * advantage * grad        return advantage  def main():    env = TuningEnv(seed=0)    agent = PolicyAgent(env.knob_sizes, lr=0.1)    print("真实最优旋钮组合:", env.optimal)    best_schedule, best_reward = None, 0.0    rewards_history = []    for step in range(400):        schedule, log_prob = agent.sample()        reward = env.reward(schedule)        if reward > best_reward:            best_reward, best_schedule = reward, schedule        agent.update(schedule, log_prob, reward)        rewards_history.append(reward)        if (step + 1) % 80 == 0:            avg = round(float(np.mean(rewards_history[-80:])), 4)            print("步数:", step + 1, "最近80步平均奖励:", avg)    print("最优发现:", best_schedule, "加速比:", round(float(best_reward), 4))    print("模拟最佳加速比: %.2fx" % best_reward)  if __name__ == "__main__":    main()

1. self.optimal = [int(rng.integers(n)) for n in self.knob_sizes]环境随机生成隐藏的最优旋钮组合，模拟真实硬件上“某个调度最快”这个未知事实。
2. time_ms = self.baseline_ms * (0.7 + 0.3 * dist)离最优越远的调度耗时越长，把抽象的性能关系简化成可计算的奖励。
3. def softmax(self, logits)把策略分数变成概率分布，是动作选择的数学基础。
4. a = int(np.random.choice(len(p), p=p))按当前策略概率采样，高概率动作更常被选中，同时保留探索机会。
5. log_prob += np.log(p[a] + 1e-12)记录被选动作的对数概率，供策略梯度计算使用。
6. advantage = reward - self.reward_mean减去历史平均奖励得到优势，让梯度更新更有针对性、方差更小。
7. grad[a] += 1.0; grad -= p计算 softmax 交叉熵梯度：推高被选动作的概率，压低其他动作。
8. self.theta[j] += self.lr * advantage * grad策略梯度更新：高优势时加强这次选择，低优势时反向修正。

常见误区

面试题