0
点赞
收藏
分享

微信扫一扫

TVM：Python 创建 Relay 表达式

code_balance 2022-06-27 阅读 5
全局变量局部变量标识符编程语言

快速了解 Relay​​¶​​

Relay IR 是纯粹的、面向表达式的语言。

从计算图的角度来看，函数（{class}​​tvm.relay.function.Function​​ ）是计算图的子图，函数调用在子图中，将其参数替换为带有相应名称的子图中的自由变量。

Relay 明显地区分了 AST 和文本格式之间的局部变量（{class}​​tvm.ir.expr.Var​​​ 使用 ​​%​​​ 标识）和全局变量（{class}​​tvm.ir.expr.GlobalVar​​​ 使用 ​​@​​ 标识）。

  • 全局标识符总是引用在全局可见环境中包含的全局可见定义，称为模块（module）。全局标识符必须是唯一的。
  • 局部标识符总是引用函数参数或被​​let​​​ （{class}​​tvm.relay.expr.Let​​​） 表达式绑定的变量，并将作用于它出现的函数或被​​let​​ 表达式绑定之处。

局部变量​​¶​​

局部变量可用于声明函数的输入参数或中间变量。可由 ​​tvm.relay.Var(name_hint, type_annotation=None)​​ 创建。其中

  • ​​name_hint​​ 指定了局部变量的名字。
  • ​​type_annotation​​ 用于局部变量的类型注解。

  In [ ]:

from tvm import relay

x = relay.Var("x") # 创建局部变量 x
x

  Out[ ]:

Var(x)

   

可以查看文本表示：

  In [ ]:

print(x)

   

free_var %x;
%x

   

如果想要声明给定 ​​dtype​​​ 和形状已知的张量的类型，可以指定 ​​type_annotation​​ 参数创建：

  In [ ]:

type_annotation = relay.TensorType(shape=(5, 5),
                                   dtype="float32")
x = relay.Var("x", type_annotation)
x

  Out[ ]:

Var(x, ty=TensorType([5, 5], float32))

   

查看文本格式：

  In [ ]:

print(x)

   

free_var %x: Tensor[(5, 5), float32];
%x

   

还有一种便捷函数：​​tvm.ir.expr.var(name_hint, type_annotation=None, shape=None, dtype="float32")​​：

   

创建变量的四种等效方式：

  In [ ]:

x = relay.Var("x", relay.TensorType([1, 2]))
x = relay.var("x", relay.TensorType([1, 2]))
x = relay.var("x", shape=[1, 2])
x = relay.var("x", shape=[1, 2], dtype="float32")

   

同样，下面两列也是等效的：

  In [ ]:

y = relay.var("x", "float32")
y = relay.var("x", shape=(), dtype="float32")

   

函数​​¶​​

Relay 中的函数的作用类似于其他编程语言中的过程或函数，并用于推广命名子图的概念。

可以直接定义函数：

  In [ ]:

a, b = [relay.var(name) for name in "ab"]
add_op = a + b
add_func = relay.Function([a, b], add_op)

   

此函数的文本形式：

  In [ ]:

add_func

  Out[ ]:

fn (%a, %b) {
  add(%a, %b)
}

   

也可以使用 Python 函数回调的形式：

  In [ ]:

def add(a, b):
    add_op = a + b
    return relay.Function([a, b], add_op)

add

  Out[ ]:

<function __main__.add(a, b)>

   

想要使用，需要：

  In [ ]:

a, b = [relay.var(name) for name in "ab"]
add_func = add(a, b)
add_func

  Out[ ]:

fn (%a, %b) {
  add(%a, %b)
}

   

也可以添加变量注解：

  In [ ]:

type_annotation = relay.TensorType(shape=(5, 5),
                                   dtype="float32")


def add(a, b):
    add_op = a + b
    return relay.Function([a, b],
                          add_op,
                          ret_type=type_annotation)


a, b = [relay.var(name, type_annotation) for name in "ab"]
add_func = add(a, b)
add_func

  Out[ ]:

fn (%a: Tensor[(5, 5), float32], %b: Tensor[(5, 5), float32]) -> Tensor[(5, 5), float32] {
  add(%a, %b)
}

   

模块​​¶​​

Relay 保留称为 “module” 的全局数据结构（在其他函数式编程语言中通常称为 “environment”），以跟踪全局函数的定义。特别地，该模块保持全局变量到它们所表示的函数表达式的全局可访问映射。模块的实用之处在于，它允许全局函数递归地引用它们自己或任何其他全局函数（例如，在 mutual 递归中）。

  In [ ]:

# 定义变量
names = "xy"
x, y = [relay.var(name) for name in names]
# 定义函数
add_op = x + y
add_func = relay.Function([x, y], add_op)

   

声明全局变量：

  In [ ]:

add_gvar = relay.GlobalVar("AddFunc")
print(add_gvar)

   

@AddFunc

   

定义将 ​​add_func​​ 提升为全局变量：

  In [ ]:

from tvm import IRModule

mod = IRModule({add_gvar: add_func})
print(mod)

   

def @AddFunc(%x, %y) {
  add(%x, %y)
}

   

获取模块的全局变量内容：

  In [ ]:

mod[add_gvar]

  Out[ ]:

fn (%x, %y) {
  add(%x, %y)
}

   

也可以直接借助全局变量的名字获取其内容：

  In [ ]:

mod["AddFunc"]

  Out[ ]:

fn (%x, %y) {
  add(%x, %y)
}

   

也可以分配新的全局变量给模块：

  In [ ]:

names = "xy"
x, y = [relay.var(name) for name in names]
# 定义函数
mul_op = x * y
mul_func = relay.Function([x, y], mul_op)
mod["MulFunc"] = mul_func

print(mod)

   

def @AddFunc(%x, %y) {
  add(%x, %y)
}

def @MulFunc(%x1, %y1) {
  multiply(%x1, %y1)
}

   

也可以通过 Python 字典更新全局变量：

  In [ ]:

names = "xyz"
x, y, z = [relay.var(name) for name in names]
# 定义函数
v1 = x * y
muladd_op = v1 + z
muladd_func = relay.Function([x, y, z], muladd_op)

mod.update({"MulAddFunc": muladd_func})
print(mod)

   

def @AddFunc(%x, %y) {
  add(%x, %y)
}

def @MulAddFunc(%x1, %y1, %z) {
  %0 = multiply(%x1, %y1);
  add(%0, %z)
}

def @MulFunc(%x2, %y2) {
  multiply(%x2, %y2)
}

   

查看全局变量：

  In [ ]:

mod.get_global_vars()

  Out[ ]:

[GlobalVar(AddFunc), GlobalVar(MulFunc), GlobalVar(MulAddFunc)]

   

回调全局变量​​¶​​

{class}​​tvm.relay.expr.Call​​ 可以在模块中回调全局变量。

比如，定义 ​​add​​ 算子：

  In [ ]:

data = relay.var("data")
bias = relay.var("bias")
add_op = data + bias

   

初始化 Relay 模块：

  In [ ]:

mod = IRModule()

   

创建并绑定 ​​add​​​ 全局函数到 ​​mod​​：

  In [ ]:

mod['AddFunc'] = relay.Function([data, bias], add_op)

   

下面定义三个变量用于定义“连加”运算：

  In [ ]:

a, b, c = [relay.var(name) for name in "abc"]

   

获取全局变量 ​​add​​：

  In [ ]:

add_gvar = mod.get_global_var('AddFunc')

   

定义“连加”运算：

  In [ ]:

add_01 = relay.Call(add_gvar, [a, b])
add_012 = relay.Call(add_gvar, [c, add_01])

   

绑定到 ​​mod​​：

  In [ ]:

mod['main'] = relay.Function([a, b, c], add_012)
print(mod)

   

def @AddFunc(%data, %bias) {
  add(%data, %bias)
}

def @main(%a, %b, %c) {
  %0 = @AddFunc(%a, %b);
  @AddFunc(%c, %0)
}

   

常量​​¶​​

{class}​​tvm.relay.Constant​​​ 节点表示常量张量值。常量被表示为 {class}​​~tvm.runtime.ndarray.NDArray​​，允许 Relay 使用 TVM 算子来计算常量。

此节点还可以表示标量常数，因为标量是形状为 ​​()​​ 的张量。因此，在文本格式中，数值和布尔字面值是将张量类型编码为零阶形状的常量的语法糖。

  In [ ]:

from tvm import nd

  In [ ]:

const_a = nd.array(7)
const_op = relay.Constant(const_a)
str(const_op), repr(const_op)

  Out[ ]:

('7i64', 'Constant(7)')

   

元组​​¶​​

元组节点构建有限（即静态已知大小）的异构数据序列。这些元组与 Python 非常匹配，它们的固定长度允许有效地投影其成员。

  In [ ]:

a = relay.var("a", shape=(10, 10))
b = relay.var("b", shape=(100, 20))
c = relay.var("c", shape=(100, 20))
f_tuple = relay.Tuple([a, b, c])

print(f_tuple)

   

free_var %a: Tensor[(10, 10), float32];
free_var %b: Tensor[(100, 20), float32];
free_var %c: Tensor[(100, 20), float32];
(%a, %b, %c)

   

支持索引：

  In [ ]:

f_tuple[0]

  Out[ ]:

Var(a, ty=TensorType([10, 10], float32))

   

If 条件表达​​¶​​

  In [ ]:

a, b = [relay.var(name) for name in "ab"]
c = a + b
d = a * b
cond = relay.const(a == b)
f = relay.If(cond, c, d)
print(f)

   

if (False) {
  free_var %a;
  free_var %b;
  add(%a, %b)
} else {
  multiply(%a, %b)
}

   

​​let​​​ 绑定​​¶​​

其实上述介绍的模块绑定属于 Relay graph 绑定，对应于计算图。

let 绑定是不可变的局部变量绑定，允许用户将表达式绑定到名称。

  • let 绑定包含局部变量、可选类型注解、值和可以引用绑定标识符的 body 表达式。如果省略了绑定变量上的类型注释，Relay 将尝试推断该变量允许的最通用类型。
  • let 表达式中的绑定变量只作用在其 body 作用域内，除非该变量定义了函数表达式。当 let 表达式创建函数时，该变量的值也在范围内，以允许递归定义函数（请参阅前一小节）。
  • let 绑定的值是计算它所依赖的绑定后的最后一个表达式的值。

  In [ ]:

x = relay.var("x")
sb = relay.ScopeBuilder()
v1 = sb.let("v1", relay.log(x))
v2 = sb.let("v2", v1 + v1)
sb.ret(v2)
f = relay.Function([x], sb.get())

f

  Out[ ]:

fn (%x) {
  let %v1 = log(%x);
  let %v2 = add(%v1, %v1);
  %v2
}

   

也可以定义 if-else 语句：

  In [ ]:

sb = relay.ScopeBuilder()
cond = relay.var("cond", 'bool')
x = relay.var("x")
y = relay.var("y")

with sb.if_scope(cond):
    one = relay.const(1, "float32")
    t1 = sb.let("t1", relay.add(x, one))
    sb.ret(t1)
with sb.else_scope():
    sb.ret(y)
f = relay.Function([x, y, cond], sb.get())
f

  Out[ ]:

fn (%x, %y, %cond: bool) {
  if (%cond) {
    let %t1 = add(%x, 1f);
    %t1
  } else {
    %y
  }
}

  更多内容见我的 Github：​​TVM 手册 — tvm-book 0.0.2 文档 (xinetzone.github.io)​​


 

探寻有趣之事！

举报

相关推荐

创建正则表达式

创建正则表达式正则表达式JavaScript前端开发

Lambda 表达式创建线程

创建线程i++OpenStack云计算

EL 表达式：${表达式}

前端jsp

python-for表达式

迭代python生成器Java编程语言

Python 条件表达式

python三元运算符运算符Linux系统/运维

Python-表达式

leetcode算法职场和发展

课时47：表达式_表达式_数字表达式

linux服务器运维bashMySQL数据库

课时45：表达式_表达式_字符串表达式

linux服务器运维字符串用户名HarmonyOS后端开发

python 表达式引擎 从条件生成数学表达式

Pythonpython字符串后端开发

C#编程：三元表达式【表达式1?表达式2:表达式3；】

程序设计.Net编程语言
0 条评论