news 2025/12/17 22:30:46

【TVM 教程】Python 目标参数化

作者头像

张小明

前端开发工程师

1.2k 24
文章封面图
【TVM 教程】Python 目标参数化

TVM 现已更新到 0.21.0 版本,TVM 中文文档已经和新版本对齐。

Apache TVM 是一个深度的深度学习编译框架,适用于 CPU、GPU 和各种机器学习加速芯片。更多 TVM 中文文档可访问 →Apache TVM

概述

对于任何支持的开发环境,TVM 都应该生成数值正确的结果。因此,在编写验证数值输出的单元测试时,这些单元测试应在所有受支持的开发环境上执行。由于这是一个非常常见的使用场景,TVM 提供了辅助函数来参数化单元测试,使它们可以在所有已启用且具有兼容设备的目标上运行。

测试套件中的一个 Python 函数可以展开成多个参数化的单元测试,每个测试针对一个单一的目标设备。要运行一个测试,必须满足以下所有条件:

单元测试文件内容

在多个目标上运行测试的推荐方法是对测试进行参数化。这可以通过使用装饰器@tvm.testing.parametrize_targets('target_1', 'target_2', ...)显式完成,并在函数中接受targetdev参数。该函数将针对列表中的每个目标运行一次,并分别报告每个目标的成功/失败情况。如果某个目标由于在 config.cmake 中被禁用,或因为没有合适的硬件而无法运行,则该目标将被标记为已跳过。

# 显式列出使用的目标 @tvm.testing.parametrize_target('llvm', 'cuda') def test_function(target, dev): # 测试代码在这里

对于需要在所有目标上正常运行的测试,可以省略装饰器。任何接受targetdev参数的测试将自动在环境变量TVM_TEST_TARGETS中指定的所有目标上参数化运行。该参数化过程会为每个目标提供相同的通过/失败/跳过报告,同时允许测试套件轻松扩展以覆盖更多目标。

# 隐式参数化运行在 TVM_TEST_TARGETS 环境变量中的所有目标上 def test_function(target, dev): # 测试代码在这里

@tvm.testing.parametrize_targets也可以用作裸装饰器来显式强调参数化,但没有额外效果。

# 显式参数化运行在 TVM_TEST_TARGETS 环境变量中的所有目标上 @tvm.testing.parametrize_targets def test_function(target, dev): # 测试代码在这里

可以使用@tvm.testing.exclude_targets@tvm.testing.known_failing_targets装饰器排除特定目标或标记预期失败的目标。有关其预期用例的更多信息,请参阅它们的文档字符串。

在某些情况下,可能需要跨多个参数进行参数化。例如,有些目标可能有多个实现方式需要测试。这种情况下,可以显式地对参数元组进行参数化,如下所示。这种写法中只会运行显式列出的目标,但每个目标仍会应用相应的@tvm.testing.requires_RUNTIME标记。

@pytest.mark.parametrize('target,impl', [ ('llvm', cpu_implementation), ('cuda', gpu_implementation_small_batch), ('cuda', gpu_implementation_large_batch), ]) def test_function(target, dev, impl): # 测试代码在这里

参数化功能是基于pytest marks实现的。每个测试函数都可以使用 pytest marks 进行装饰以添加元数据。最常用的标记如下:

在使用目标参数化时,每次测试运行都会被装饰上与其目标对应的@tvm.testing.requires_RUNTIME。因此,如果目标在config.cmake中被禁用或没有可用硬件,将会被明确标记为跳过。

还存在一个tvm.testing.enabled_targets()函数,它会根据环境变量TVM_TEST_TARGETS、构建配置以及实际硬件返回所有已启用并可运行的目标。目前多数测试都显式地循环遍历该函数的返回值,但这不应用于新测试。这种写法在 pytest 输出中会悄悄跳过禁用的运行时,或无法运行的设备,而且测试一旦在某个目标失败就会中止运行,使得难以判断是该目标出错还是所有目标都失败。

# 旧式写法,不推荐使用 def test_function(): for target,dev in tvm.testing.enabled_targets(): # 测试代码在这里

本地运行

在本地运行 Python 单元测试,可以在${TVM_HOME}目录下使用命令pytest

注意:如果TVM_TEST_TARGETS中不包含任何既被启用又有可用设备的目标,则测试将回退,仅在llvm目标上运行。

注意:此过滤是在基于环境变量TVM_TEST_TARGETS选择目标之后执行的。即使指定了-m gpu,如果TVM_TEST_TARGETS中不包含 GPU 目标,也不会运行 GPU 测试。

在本地 Docker 容器中运行

可以使用docker/bash.sh脚本在与 CI 使用的相同 Docker 镜像中运行单元测试。第一个参数应指定要运行的 Docker 镜像(例如docker/bash.sh ci_gpu)。允许的镜像名称在 TVM 源代码目录中的 Jenkinsfile 文件顶部定义,并映射到 tlcpack 上的镜像。

如果不提供额外参数,Docker 镜像将启动一个交互式 bash 会话。如果传入脚本作为可选参数(如docker/bash.sh ci_gpu tests/scripts/task_python_unittest.sh),该脚本将在 Docker 镜像中被执行。

注意:Docker 镜像包含所有系统依赖项,但不包含该系统的build/config.cmake配置文件。TVM 源目录将作为 Docker 镜像的主目录,因此默认会使用本地的 config/build 目录。一个可行的做法是分别维护build_localbuild_docker目录,在进入或退出 Docker 时将build符号链接到相应目录。

在 CI 中运行

CI 中的所有流程都起始于 Jenkinsfile 中定义的任务。这些定义包括指定使用的 Docker 镜像、编译时配置和每个阶段所运行的测试。

Jenkinsfile 中的每个任务(如“BUILD: CPU”)都会调用docker/bash.sh。紧随其后的参数定义了 CI 中使用的 Docker 镜像,方式与本地执行一致。

Docker 镜像中并不包含config.cmake文件,因此这是每个BUILD任务的第一步。此步骤通过执行tests/scripts/task_config_build_*.sh脚本完成。具体使用哪个脚本取决于正在测试的构建方式,并在 Jenkinsfile 中指定。

每个BUILD任务最终都会打包一个库文件,供后续测试阶段使用。

Jenkinsfile 中的Unit TestIntegration Test阶段定义了如何调用pytest。每个任务都以解压先前在BUILD阶段编译的库文件开始,随后运行一个测试脚本(如tests/script/task_python_unittest.sh)。这些脚本会设置要传递给pytest的文件/目录和命令行参数。

多个测试脚本使用了-m gpu选项,用于限定仅运行带有@pytest.mark.gpu标记的测试用例。

版权声明: 本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若内容造成侵权/违法违规/事实不符,请联系邮箱:809451989@qq.com进行投诉反馈,一经查实,立即删除!
网站建设 2025/12/17 22:29:32

收藏!大模型时代必看:6大热门岗位,小白程序员转岗指南

当下,大模型技术正以“破竹之势”渗透到各行各业,从智能客服到自动驾驶,从代码生成到科学研究,处处都能看到它的身影。随之而来的,是大模型相关岗位的爆发式增长,成为程序员和技术小白转型的黄金赛道。 今天…

作者头像 李华
网站建设 2025/12/17 22:29:22

内网渗透学习必备干货:12 个高价值靶场平台 + 实战避坑指南 + 操作技巧全收录!

前言 在内网渗透学习中,“实战靶场” 是连接理论与实操的核心桥梁 —— 它能模拟真实企业内网的拓扑结构、漏洞分布和信任关系,让你在合法可控的环境中练手 “跳板机横向移动”“域控提权” 等关键技能。 以下按 “新手入门→进阶实战→专项突破” 三个…

作者头像 李华
网站建设 2025/12/17 22:24:57

你不知道的Q#调试黑科技:如何让Python端实时捕获量子状态异常

第一章:你不知道的Q#调试黑科技:如何让Python端实时捕获量子状态异常在混合量子-经典计算架构中,Q# 与 Python 的协同调试长期面临“黑盒”困境——量子态无法直接观测,异常往往滞后暴露。然而,借助 Q# 的 DumpMachine…

作者头像 李华
网站建设 2025/12/17 22:24:51

新型僵尸网络正对路由器、摄像头等设备发起大规模DDoS攻击

新型僵尸网络正对路由器、摄像头等设备发起大规模DDoS攻击 趋势科技的一项新研究发现,自 2024 年底以来,一个新发现的物联网 (IoT) 僵尸网络一直利用路由器、IP 摄像头和其他连接设备等物联网设备中的漏洞,在全球策划…

作者头像 李华