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_width和pad_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.]]]]