nn.pad

nn.pad#

op.nn.pad 函数,用于对输入数据进行填充。下面是对该函数的解读:

该函数接受以下参数:

  • data:输入数据的表达式(tvm.relay.Expr类型)

  • pad_width:每个轴要填充的宽度,以元组的形式给出,格式为 ((before_1, after_1), ..., (before_N, after_N))

  • pad_value:(可选)填充的值,默认为 0

  • pad_mode:(可选)填充模式,可以是 'constant''edge''reflect',分别表示使用常量值、边缘值或反射值进行填充

函数首先检查pad_widthpad_value的类型,如果它们不是预期的类型,则进行相应的转换。然后根据 pad_width 的类型选择不同的填充方式,并返回计算结果。

import numpy as np
import tvm
from tvm import relay
from tvm.relay import op
from tvm.relay.testing import run_opt_pass
dshape = 1, 2, 1, 1
pad_width = [(1, 0), (0, 1), (0, 0), (0, 0)]
x = relay.var("x", shape=dshape)
y = op.nn.pad(x, pad_width, pad_value=-1, pad_mode='constant')
pad_width[:2] = [(-1, 0), (0, -1)]
x_ = op.nn.pad(y, pad_width, pad_value=-1, pad_mode='constant') # 移除 pad
t = relay.Tuple([y, x_])
func = relay.Function([x], t)
func = run_opt_pass(func, relay.transform.InferType())
tvm.IRModule.from_expr(func).show()
intrp = relay.create_executor("graph", device=tvm.cpu(0), target="llvm")

data_np = np.arange(np.prod(dshape)).reshape(dshape).astype("float32")
print(f"原始数据({data_np.shape}): \n{data_np}")
op_res, new_data = intrp.evaluate(func)(data_np)
print(f"最终数据({op_res.shape}): \n{op_res}")
np.testing.assert_allclose(data_np, new_data.numpy())
def @main(%x: Tensor[(1, 2, 1, 1), float32] /* ty=Tensor[(1, 2, 1, 1), float32] */) -> (Tensor[(2, 3, 1, 1), float32], Tensor[(1, 2, 1, 1), float32]) {
  %0 = nn.pad(%x, -1 /* ty=int32 */, pad_width=[[1, 0], [0, 1], [0, 0], [0, 0]]) /* ty=Tensor[(2, 3, 1, 1), float32] */;
  %1 = nn.pad(%0, -1 /* ty=int32 */, pad_width=[[-1, 0], [0, -1], [0, 0], [0, 0]]) /* ty=Tensor[(1, 2, 1, 1), float32] */;
  (%0, %1) /* ty=(Tensor[(2, 3, 1, 1), float32], Tensor[(1, 2, 1, 1), float32]) */
}
原始数据((1, 2, 1, 1)): 
[[[[0.]]

  [[1.]]]]
最终数据((2, 3, 1, 1)): 
[[[[-1.]]

  [[-1.]]

  [[-1.]]]


 [[[ 0.]]

  [[ 1.]]

  [[-1.]]]]