趋近智
优化深度神经网络 (neural network)需要从理论搜索空间转到具体实现。自动化搜索和成本建模的原理应用于标准ResNet瓶颈块。我们使用现代的基于生成的搜索方法,通常称为自动调度(特别是TVM中的Ansor系统),而不是基于模板的自动调优。此方法使编译器能够自动生成涉及算子融合的复杂子图的搜索空间,而不是依赖手动编写的调度模板。
一个ResNet块包含多个卷积层、批归一化 (normalization)和激活函数 (activation function),以及一个残差加法。当将其降级到硬件目标时,编译器将其视为一个子图。我们的目标是为整个子图确定最佳循环调度,让编译器能够进行积极的算子融合和分块。
我们首先使用高级中间表示(Relay)定义网络布局。虽然我们着重于ResNet-50的典型形状(输入:),但此方法适用于任何张量计算。
import tvm
from tvm import relay, auto_scheduler
import numpy as np
def get_resnet_block(batch_size=1, dtype="float32"):
data_shape = (batch_size, 64, 56, 56)
data = relay.var("data", shape=data_shape, dtype=dtype)
weight1 = relay.var("weight1", shape=(64, 64, 3, 3), dtype=dtype)
# 卷积 -> 批归一化 -> ReLU
conv1 = relay.nn.conv2d(data, weight1, padding=(1, 1), kernel_size=(3, 3))
bn1 = relay.nn.batch_norm(conv1, relay.var("gamma"), relay.var("beta"),
relay.var("mean"), relay.var("var"))[0]
relu1 = relay.nn.relu(bn1)
# 模拟残差加法(为演示简化)
out = relay.add(relu1, data)
return relay.Function(relay.analysis.free_vars(out), out)
搜索过程包含一个迭代反馈循环。搜索策略从上一节定义的搜索空间生成一批候选程序(草图)。这些候选程序在硬件上进行测量,结果用于更新统计成本模型(XGBoost或LightGBM)。
这个调优循环的架构推动优化。TaskScheduler协调整个过程,如果我们要调优完整网络,它会为不同的子图分配时间资源。对于单个块,它严格专注于最大化我们定义函数的吞吐量 (throughput)。
自动调度流程。搜索策略生成候选程序,这些程序在进行实际基准测试之前会通过成本模型进行筛选。
为了执行调优,我们从Relay函数中提取搜索任务。编译器分析图以识别可优化的子图。然后我们配置 Tuner。
在生产环境中,measure_ctx 通常会使用RPC跟踪器指向远程设备。对于这个本地示例,我们假设主机是目标。
# 1. 从网络中提取任务
target = tvm.target.Target("llvm -mcpu=skylake")
mod = tvm.IRModule.from_expr(get_resnet_block())
tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], [], target)
# 2. 配置搜索选项
log_file = "resnet_block_tuning.json"
measure_ctx = auto_scheduler.LocalRPCMeasureContext(min_repeat_ms=300)
tuner = auto_scheduler.TaskScheduler(tasks, task_weights)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=1000, # 要测试的配置总数
runner=measure_ctx.runner,
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
verbose=1 # 设置为 0 表示静默运行
)
# 3. 执行搜索
# 此过程可能需要几分钟到几小时,耗时视硬件而定
tuner.tune(tune_option)
执行期间,调优器搜索由循环分块大小 、向量 (vector)化因子和展开深度定义的空间 。最初,由于演化搜索从低效的默认调度中移开,性能改进会很快。随着过程持续,增益变得渐进,表现为汇编生成中的微优化。
成本模型在此很重要。没有它,调优器将依赖随机搜索。该模型学习特定循环结构(特征)与执行时间(标签)之间的相关性,使其能够有效地修剪搜索空间,而无需在硬件上运行每个候选程序。
调优完成后,我们分析改进情况。下图显示了卷积核的典型收敛模式。X轴表示测量试验次数,Y轴表示GFLOPS的吞吐量 (throughput)。
自动调优会话期间的性能收敛。 “最佳结果”线记录了截至该试验所识别的最高吞吐量配置。
在最初的100次试验中,调优器通常会发现最佳分块因子(例如,拆分循环以适应L1/L2缓存)。后续的改进通常源于向量 (vector)化宽度调整和解决线程绑定冲突。
调优过程的输出是一个JSON日志文件,包含找到的最佳调度的参数 (parameter)。要在生产推理 (inference)流程中使用此文件,我们使用 ApplyHistoryBest 编译模型。此通道用我们在搜索期间找到的特定调度替换默认的降级机制。
# 使用历史最佳上下文编译
with auto_scheduler.ApplyHistoryBest(log_file):
with tvm.transform.PassContext(opt_level=3, config={"relay.backend.use_auto_scheduler": True}):
lib = relay.build(mod, target=target, params=None)
# 创建运行时模块
dev = tvm.device(str(target), 0)
module = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
# 基准测试
print("优化后的推理已就绪。")
当 ApplyHistoryBest 激活时,编译器会将Relay图中的工作负载签名与日志文件中的条目进行匹配。如果找到匹配项,则应用相应的低级IR (TIR) 转换。如果没有找到匹配项(例如,如果输入形状更改),编译器将回退到默认的、未优化的调度,这表明针对部署中遇到的特定形状进行调优很有用。
工程师经常遇到调优器未能找到比基线更好的调度,或搜索崩溃的情况。常见原因包括:
为了减少这些问题,请确保 target 字符串准确反映处理器能力(例如,包含AVX-512等特定向量扩展或张量核心版本)。此外,使用 TransferLearning 会有帮助;您可以从类似的硬件目标加载日志文件来初始化成本模型,从而给搜索一个“热启动”。
这部分内容有帮助吗?
© 2026 ApX Machine LearningAI伦理与透明度•