tvm.relax.frontend

用于构建 Relax 程序的前端，以及模型导入器

tvm.relax.frontend.detach_params(mod: IRModule) → Tuple[IRModule, Dict[str, List[NDArray]]]

将输入 IRModule 函数中的 “params” 属性分离为单独的参数字典。

参数:

mod (tvm.IRModule) – 要分离其函数 “param” 属性的 IRModule。

返回:

detached_mod (tvm.IRModule) – 分离后的 IRModule。
params_dict (Dict[str, List[tvm.nd.NDArray]]) – 分离的参数。字典键对应于输入 IRModule 中具有 “params” 属性的函数的名称。

tvm.relax.frontend.nn

一个类似 PyTorch 的 API，用于构建 IRModule。

class tvm.relax.frontend.nn.Effect

Effect 是一种特殊的非用户界面类型，用于表示具有副作用的操作，例如 print。它用于表示计算的输出。

emit_init(name_hint: str, builder: BlockBuilder) → List[DataflowVar]: 发射 effect 的初始化。此方法由编译器调用以初始化 effect。

create(name_hint: str) → List[Var]: 创建 relax.Function 的隐式输入，该函数表示副作用

set_state(state_vars: List[Var]) → None: 设置表示 effect 的变量

finalize() → List[Var]: 将 effect 最终化为 relax.Function 的隐式返回值

to(dtype: str | None = None) → None: 将 effect 转换为特定 dtype。通常对于大多数 effect 来说，它是无操作的。

class tvm.relax.frontend.nn.Module

神经网络组件的基类。子类化它可以构建您的模型。模块可以使用常规属性赋值在树结构中相互嵌套。

named_parameters(prefix: str = '') → Iterator[Tuple[str, Parameter]]

此方法提供了一个模块参数的迭代器，同时生成参数名称及其对应的值。

参数:: prefix (str) – 要添加到所有参数名称的前缀。
生成:: (str, Parameter) - 包含名称和参数的元组

parameters() → Iterator[Parameter]

此方法提供了一个模块参数的迭代器，仅生成 Parameter 值。

生成:: Parameter - 模块的参数

state_dict(*, prefix: str = '', destination: Dict[str, Parameter] | None = None) → Dict[str, Parameter]

返回一个字典，其中包含对模块整个状态的引用。

参数:

prefix (str) – 要添加到所有参数名称的前缀。
destination (Optional[Dict[str, Parameter]]) – 将状态保存到的字典。如果为 None，则创建一个新字典。

返回:

dict – 包含模块整个状态的字典

返回类型:

Dict[str, Parameter]

load_state_dict(state_dict: Dict[str, Parameter], strict: bool = True) → Tuple[List[str], List[str]]

此函数将参数和缓冲区从 state_dict 复制到当前模块及其后代。如果 strict 设置为 True，则 state_dict 中的键必须与此模块的 state_dict() 函数返回的键完全匹配。

参数:

state_dict (Dict[str, Parameter]) – 包含模块整个状态的字典
strict (bool = True) – 是否严格强制 state_dict 中的键与此模块的 state_dict() 函数返回的键匹配。

返回:

(missing_keys, unexpected_keys) – 包含两个列表的元组：缺失的键和意外的键。

返回类型:

Tuple[List[str], List[str]]

to(dtype: str | None = None) → None: 递归地将模块转换为特定 dtype

export_tvm(spec: _spec.ModuleSpecType, debug: bool = False, allow_extern: bool = False) → Tuple[IRModule, List[Tuple[str, Parameter]]] | Tuple[IRModule, List[Tuple[str, Parameter]], List[ExternModule]]

将模块导出到 TVM IRModule 和参数

参数:

spec (_spec.ModuleSpecType) – 一个字典，将每个输入名称映射到一个规范，该规范定义了输入形状和 dtype。
debug (bool) – 如果设置为 True，则导出的模块将支持 effects。这使得在图中进行打印等操作成为可能。

返回:

irmodule (tvm.ir.IRModule) – 模型转换后的 tvm IR 表示形式。
params (List[Tuple[str, Parameter]]) – 与模型权重相对应的参数列表。
ext_mods (List[nn.ExternModule]) – 模型中使用的 ExternModule 列表。

jit(spec: _spec.ModuleSpec, device: str | Device = 'cpu', pipeline: None | str | Pass = 'default_build', out_format: str = 'torch', debug: bool = False) → Any: 将 nn.model 即时编译为可执行文件

class tvm.relax.frontend.nn.ModuleList(modules: List[Module])

在列表中保存子模块。

append(module: Module): 将模块添加到 ModuleList 的末尾

to(dtype: str | None = None) → None: 递归地将模块转换为特定 dtype

forward(x): 模块的前馈传递

class tvm.relax.frontend.nn.Object(*, _expr: RelaxExpr, _name: str): relax.Expr 之上的包装器，其 struct_info 是基本 ObjectStructInfo（而不是其任何子类）。Object 有效地表示非张量前端组件，例如 KV 缓存。

class tvm.relax.frontend.nn.Parameter(shape: Sequence[int | str | PrimExpr], dtype: str | None = None)

参数表示神经网络层的权重。它是一个特殊的张量，可以绑定或不绑定到具体值。如果参数绑定到具体值，则称为绑定参数，否则称为未绑定参数。

property data: NDArray | None: 如果参数绑定到具体值，则返回参数的具体值，否则返回 None。返回的值是 tvm.runtime.NDArray。

to(dtype: str | None = None) → None: 如果参数未绑定到任何具体数据，则更改参数的 dtype。

class tvm.relax.frontend.nn.Tensor(*, _expr: RelaxExpr)

relax.Expr 之上的包装器，其 struct_info 是 TensorStructInfo，提供更方便的访问形状和 dtype 信息。Tensor 始终是符号化的，并且不绑定到任何具体值。形状和 dtype 推理在张量创建时立即完成，即当算子应用于张量时，形状和 dtype 信息已可用。

static from_const(data) → Tensor: 从 numpy 常量构造张量。

static from_scalar(data: int | float, dtype: str) → Tensor: 从标量构造张量，并指定 dtype。

static from_struct_info(struct_info: TensorStructInfo, name: str = 'tensor') → Tensor: 从 relax TensorStructInfo 构造 nn.Tensor

static placeholder(shape: Sequence[int | str | PrimExpr], dtype: str, name: str = 'tensor') → Tensor

创建具有给定形状和 dtype 的占位符张量。在通常情况下，用户永远不应直接创建占位符张量，唯一的例外是指示外部函数的返回值形状/dtype。

如果 shape 是字符串名称，我们创建一个符号形状 tvm.tir.Var(name, “int64”)。

property shape: List[int | PrimExpr]

以整数列表形式返回张量的形状。

整数可以是 python int 或 tvm.tir.PrimExpr，取决于形状是否完全静态，例如，[1, 2, tvm.tir.Var(“n”)] 是一个有效的形状，其中最后一个维度是动态的，而前两个维度始终是静态常量。

返回:: shape – 张量的形状
返回类型:: List[Union[int, tir.PrimExpr]]

property ndim: int

返回张量的维度数量。

返回:: ndim – 张量的维度数量
返回类型:: int

property dtype: str

返回张量的数据类型。

返回:: dtype – 张量的数据类型
返回类型:: str

tvm.relax.frontend.nn.add_extern(mod: ExternModule) → None: 向导出器添加一个外部模块。

class tvm.relax.frontend.nn.ExternModule(symbols: Dict[str, Callable])

外部模块的抽象基类。外部模块旨在帮助将用户手工编写的内核合并到导出的 TVM IRModule 中。

load() → Module: 将外部模块加载到 TVM 运行时模块中。

class tvm.relax.frontend.nn.ObjectModule(symbols: Dict[str, Callable], filepath: Path)

nn.ExternModule 的子类，允许用户提供要链接到编译产物中的对象 .o 文件；

load() → Module: 将外部模块加载到 TVM 运行时模块中。

class tvm.relax.frontend.nn.SourceModule(symbols: Dict[str, Callable], source_code: str | Path, source_format: str, compile_options: List[str] | None = None, compiler: str | None = None, output_format: str = 'obj')

nn.ExternModule 的子类。它编译 C++/CUDA 源代码并将它们链接到最终的 IRModule 中。

形状/数据类型推断。 nn.ExternModule 系统要求用户提供额外的信息才能工作，即 symbols。它是一个字典，将外部目标文件中的每个符号映射到其形状/数据类型推断函数。考虑一个函数 my_func 接受两个张量的情况，形状为 (x, y, 1) 的 a 和形状为 (y, z, 5) 的 b，并生成形状为 (x, y, z, 9) 的张量 c，形状/数据类型推断函数应如下所示

def shape_dtype_inference(a, b):
    x, y, _ = a.shape
    _, z, _ = b.shape
    return nn.Tensor.placeholder((x, y, z, 9), dtype="float32")

并且 symbols 字典应提供为

symbols={
    "my_func": shape_dtype_inference,
}

调用约定。 所有外部模块现在都遵循 “目标传递风格”（DPS）调用约定，这意味着返回的张量已由系统预先分配，并作为外部函数的参数传入。

重用上面的例子，my_func 的实现应在其签名中包含三个参数，其中张量使用来自 DLPack 的 DLTensor 表示，DLPack 是张量内存表示的事实标准。更多详情：https://github.com/dmlc/dlpack/blob/v0.8/include/dlpack/dlpack.h#L163-L206。

为了暴露符号，保证 TVM_DLL_EXPORT_TYPED_FUNC(symbol, function) 可用

// those headers are guaranteed to be available
#include <dlpack/dlpack.h>
#include <tvm/runtime/data_type.h>
#include <tvm/runtime/packed_func.h>

namespace {
// anonymous namespace hides the symbol `_my_func_impl` from other translation units
int _my_func_impl(DLTensor* a, DLTensor* b, DLTensor* c) {
    // `a` and `b` are inputs, and `c` is the output
}
}
// expose symbol `my_func` instead of `_my_func_impl`
TVM_DLL_EXPORT_TYPED_FUNC(my_func, _my_func_impl);

编译器Pass `AttachExternModules`。 引入它是为了在编译管道的任何阶段将 nn.ExternModule 列表附加到 IRModule 中，并将编译后的外部模块作为 `runtime.Module` 附加到 IRModule 的 `external_mods` 属性中。 tvm.compile 链接需要它，但有了这个 pass 的存在，源代码编译可以推迟到 TVM 编译的任意阶段。

注意事项。 需要在 export_tvm 期间精确调用一次 nn.add_extern 来注册外部模块。每个符号应该只注册一次，以避免潜在的冲突，否则会引发错误。

static tvm_home() → Path

查找 TVM 的 home 目录。如果设置了 TVM_HOME 环境变量，则使用它。否则，使用安装 tvm Python 包的目录。作为健全性检查，要求具有 include 和 3rdparty 作为直接子目录。

返回:: tvm_home – TVM 的 home 目录，并保证具有 include 和 3rdparty 作为直接子目录。
返回类型:: pathlib.Path

static get_includes(tvm_pkg: List[str] | None = None) → List[Path]

根据 tvm_home() 返回默认的 include 路径。默认情况下，它包括 TVM、DLPack 和 DMLC-Core。如果提供了 tvm_pkg，它还包括 tvm_home/3rdparty 下指定的软件包。

参数:: tvm_pkg (可选[List[str]]) – 要包含在 tvm_home/3rdparty 下的软件包列表。每个元素都应该是 tvm_home/3rdparty 的相对路径。
返回:: includes – include 路径列表。
返回类型:: List[pathlib.Path]

static get_compile_options(source_format: str, tvm_pkg: List[str] | None = None) → List[str]

根据 source_format 返回默认的编译选项，包括相对于 tvm_home() 的默认 include 路径、配置 DMLC-Core 的默认标志，并且默认情况下，它使用 “-O3” 和 “-std=c++17”。

参数:

source_format (str) – 源代码格式。它可以是 “cpp” 或 “cu”。
tvm_pkg (可选[List[str]]) – 要包含在 tvm_home/3rdparty 下的软件包列表。每个元素都应该是 tvm_home/3rdparty 的相对路径。

返回:

compile_options – 编译标志列表。

返回类型:

List[str]

compile(output_path: Path) → None: 在提供的目录中编译源代码并返回编译后的产物。

load() → Module: 将外部模块加载到 TVM 运行时模块中。

class tvm.relax.frontend.nn.GELU: 用于 GELU 激活层的 relax.frontend.nn.Module。

class tvm.relax.frontend.nn.Conv1D(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: str | None = None)

用于 conv1d 层的 relax.frontend.nn.Module。

forward(x: Tensor) → Tensor

conv1d 层的前向方法。

参数:: x (Tensor) – 输入张量。
返回:: ret – conv1d 层的输出张量。
返回类型:: Tensor

class tvm.relax.frontend.nn.Conv2D(in_channels: int, out_channels: int, kernel_size: List[int] | int, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: str | None = None, data_layout: str = 'NCHW')

用于 conv2d 层的 relax.frontend.nn.Module。

forward(x: Tensor) → Tensor

conv2d 层的前向方法。

参数:: x (Tensor) – 输入张量。
返回:: ret – conv2d 层的输出张量。
返回类型:: Tensor

class tvm.relax.frontend.nn.Conv3D(in_channels: int, out_channels: int, kernel_size: List[int] | int, stride: List[int] | int = 1, padding: List[int] | int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: str | None = None, data_layout: str = 'NCDHW')

用于 conv3d 层的 relax.frontend.nn.Module。

forward(x: Tensor) → Tensor

conv3d 层的前向方法。

参数:: x (Tensor) – 输入张量。
返回:: ret – conv3d 层的输出张量。
返回类型:: Tensor

class tvm.relax.frontend.nn.ConvTranspose1D(in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, output_padding: int = 0, dilation: int = 1, groups: int = 1, bias: bool = True, dtype: str | None = None)

用于 ConvTranspose1D 层的 relax.frontend.nn.Module。

forward(x: Tensor) → Tensor

conv transpose 1d 层的前向方法。

参数:: x (Tensor) – 输入张量。
返回:: ret – conv transpose 1d 层的输出张量。
返回类型:: Tensor

用于 embedding 层的 relax.frontend.nn.Module。

forward(x: Tensor)

embedding 层的前向方法。

参数:: x (Tensor) – 输入张量。
返回:: ret – embedding 层的输出张量。
返回类型:: Tensor

class tvm.relax.frontend.nn.GroupNorm(num_groups: int, num_channels: int, eps: float = 1e-05, affine: bool = True, dtype: str | None = None)

用于 group norm 层的 relax.frontend.nn.Module。

forward(x: Tensor, channel_axis: int = 1, axes: List[int] | None = None)

group norm 层的前向方法。

参数:

x (Tensor) – 输入张量。
channel_axis (int) – 输入数据的通道轴。
axes (Optional[List[int]]) – 用于计算范数的轴的可选列表，如果未指定，则假定前两个轴应保持不变。

返回:

ret – group norm 层的输出张量。

返回类型:

tvm.relax.frontend

tvm.relax.frontend.nn

tvm.relax.frontend.onnx

tvm.relax.frontend.stablehlo

tvm.relax.frontend.torch