tvm.relax.testing.nn

tvm.relax.testing.nn#

import tvm
from tvm import relax
from tvm.relax.testing import nn
from tvm.script import ir as I, relax as R, tir as T

emit#

class ReLU(nn.Module):
    def forward(self, x: relax.Expr) -> relax.Var:
        return nn.emit(R.nn.relu(x), name_hint="x")

bb = relax.BlockBuilder()
with bb.function("main"):
    model = ReLU()
    x = nn.Placeholder((32, 32), dtype="float32", name="x")
    output = model(x)
    params = [x] + model.parameters()
    bb.emit_func_output(output, params)
mod = bb.get()

mod.show()

# from tvm.script import ir as I
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function
    def main(x: R.Tensor((32, 32), dtype="float32")) -> R.Tensor((32, 32), dtype="float32"):
        x1: R.Tensor((32, 32), dtype="float32") = R.nn.relu(x)
        return x1

参数#

class Plus1(nn.Module):
    def __init__(self):
        self.const_1 = relax.const(1, "float32")

    def forward(self, x: relax.Expr) -> relax.Var:
        return nn.emit(R.add(x, self.const_1))

model = Plus1()
assert model.parameters() == []

tvm.realx.testing.nn.Module.define_subroutine#

当 nn.Module 类的 define_subroutine 属性设置为 True 时，模块会被正确定义为子例程（subroutine），而不是内联到调用者代码中。

class Activation(nn.Module):
    define_subroutine = True

    def forward(self, state: relax.Expr) -> relax.Var:
        return R.nn.relu(state)

class Layer(nn.Module):
    define_subroutine = True

    def __init__(self, in_features, out_features):
        self.weights = nn.Parameter(
            (in_features, out_features), dtype="float32", name="weights"
        )
        self.activation = Activation()

    def forward(self, x: relax.Expr) -> relax.Var:
        state = R.matmul(x, self.weights)
        return self.activation(state)

  • Activation 类：简单的 ReLU 激活函数层，设置 define_subroutine = True

  • Layer 类：包含权重参数和激活函数的神经网络层，同样设置 define_subroutine = True

通过 define_subroutine = True 设置，确保：

  1. Activation 和 Layer 类被编译为独立的子例程函数

  2. 生成的 IR 模块中包含私有函数 layer 和 activation

  3. 主函数 main 正确调用这些子例程

该功能允许开发者将复杂的模型分解为可重用的子例程，提高代码的模块化和可维护性，同时有助于 TVM 优化器进行更有效的优化。

model = Layer(64, 32)
batch_size = tvm.tir.Var("batch_size", "int64")
x = nn.Placeholder((batch_size, 64), dtype="float32", name="x")

bb = relax.BlockBuilder()
with bb.function("main", params=[x, *model.parameters()]):
    output = model(x)
    bb.emit_func_output(output)
mod = bb.get()
mod.show()

# from tvm.script import ir as I
# from tvm.script import tir as T
# from tvm.script import relax as R

@I.ir_module
class Module:
    @R.function(private=True)
    def activation(state: R.Tensor(("batch_size", 32), dtype="float32")) -> R.Tensor(("batch_size", 32), dtype="float32"):
        batch_size = T.int64()
        gv1: R.Tensor((batch_size, 32), dtype="float32") = R.nn.relu(state)
        return gv1

    @R.function(private=True)
    def layer(x: R.Tensor(("batch_size", 64), dtype="float32"), weights: R.Tensor((64, 32), dtype="float32")) -> R.Tensor(("batch_size", 32), dtype="float32"):
        batch_size = T.int64()
        cls = Module
        gv: R.Tensor((batch_size, 32), dtype="float32") = R.matmul(x, weights, out_dtype="void")
        gv2: R.Tensor((batch_size, 32), dtype="float32") = cls.activation(gv)
        return gv2

    @R.function
    def main(x: R.Tensor(("batch_size", 64), dtype="float32"), weights: R.Tensor((64, 32), dtype="float32")) -> R.Tensor(("batch_size", 32), dtype="float32"):
        batch_size = T.int64()
        cls = Module
        gv3: R.Tensor((batch_size, 32), dtype="float32") = cls.layer(x, weights)
        return gv3
目录