如何写新的Python OP

    py_func 具体接口为：

    其中，

    • x 是Python Op的输入变量，可以是单个 Tensor | tuple[Tensor] | list[Tensor] 。多个Tensor以tuple[Tensor]或list[Tensor]的形式传入。

    • out 是Python Op的输出变量，可以是单个 Tensor | tuple[Tensor] | list[Tensor]，也可以是Numpy Array 。

    • func 是Python Op的前向函数。在运行网络前向时，框架会调用 out = func(*x) ，根据前向输入 x 和前向函数 func 计算前向输出 out。在 func 建议先主动将Tensor转换为numpy数组，方便灵活的使用numpy相关的操作，如果未转换成numpy，则可能某些操作无法兼容。

    • backward_func 是Python Op的反向函数。若 backward_func 为 None ，则该Python Op没有反向计算逻辑； 若 backward_func 不为 None，则框架会在运行网路反向时调用 backward_func 计算前向输入 x 的梯度。

    • skip_vars_in_backward_input 为反向函数 backward_func 中不需要的输入，可以是单个 Tensor | tuple[Tensor] | list[Tensor] 。

    • 第一步：定义前向函数和反向函数

    前向函数和反向函数均由Python编写，可以方便地使用Python与numpy中的相关API来实现一个自定义的OP。

    若前向函数的输入为 x_1, x_2, …, x_n ，输出为y_1, y_2, …, ，则前向函数的定义格式为：

    1. def foward_func(x_1, x_2, ..., x_n):
    2.     ...
    3.     return y_1, y_2, ..., y_m

    默认情况下，反向函数的输入参数顺序为：所有前向输入变量 + 所有前向输出变量 + 所有前向输出变量的梯度，因此对应的反向函数的定义格式为：

    若反向函数不需要某些前向输入变量或前向输出变量，可设置 skip_vars_in_backward_input 进行排除（步骤三中会叙述具体的排除方法）。

    注：，x_1, …, x_n为输入的多个Tensor，请以tuple(Tensor)或list[Tensor]的形式在py_func中传入。建议先主动将Tensor通过numpy.array转换为数组，否则Python与numpy中的某些操作可能无法兼容使用在Tensor上。

    此处我们利用numpy的相关API完成tanh的前向函数和反向函数编写。下面给出多个前向与反向函数定义的示例：

    1. import numpy as np
    3. # 前向函数1：模拟tanh激活函数
    4. def tanh(x):
    5.     # 可以直接将Tensor作为np.tanh的输入参数
    6.     return np.tanh(x)
    8. # 前向函数2：将两个2-D Tenosr相加，输入多个Tensor以list[Tensor]或tuple(Tensor)形式
    9.     # 必须先手动将Tensor转换为numpy数组，否则无法支持numpy的shape操作
    10.     x = np.array(x)  
    11.     y = np.array(y)
    13.     if x.shape != y.shape:
    14.         raise AssertionError("the shape of inputs must be the same!")
    16.     result = np.zeros(x.shape, dtype='int32')
    17.     for i in range(len(x)):
    18.         for j in range(len(x[0])):
    19.             result[i][j] = x[i][j] + y[i][j]
    21.     return result
    23. # 前向函数3：可用于调试正在运行的网络（打印值）
    24. def debug_func(x):
    25.     # 可以直接将Tensor作为print的输入参数
    26.     print(x)
    28. # 前向函数1对应的反向函数，默认的输入顺序为：x、out、out的梯度
    30.     return np.array(dy) * (1 - np.square(np.array(y)))

    注意，前向函数和反向函数的输入均是 Tensor 类型，输出可以是Numpy Array或 Tensor。 由于 Tensor 实现了Python的buffer protocol协议，因此即可通过 numpy.array 直接将 Tensor 转换为numpy Array来进行操作，也可直接将 Tensor 作为numpy函数的输入参数。但建议先主动转换为numpy Array，则可以任意的使用python与numpy中的所有操作（例如”numpy array的+/-/shape”）。

    • 第二步：创建前向输出变量

    我们需调用 Program.current_block().create_var 创建前向输出变量。在创建前向输出变量时，必须指明变量的名称name、数据类型dtype和维度shape。

    1. import paddle
    3. paddle.enable_static()
    5. def create_tmp_var(program, name, dtype, shape):
    6.     return program.current_block().create_var(name=name, dtype=dtype, shape=shape)
    8. in_var = paddle.static.data(name='input', dtype='float32', shape=[-1, 28, 28])
    10. # 手动创建前向输出变量
    11. out_var = create_tmp_var(paddle.static.default_main_program(), name='output', dtype='float32', shape=[-1, 28, 28])
    • 第三步：调用 py_func 组建网络

    py_func 的调用方式为：

    若我们不希望在反向函数输入参数中出现前向输入，则可使用 skip_vars_in_backward_input 进行排查，简化反向函数的参数列表。

    1. paddle.static.nn.py_func(func=tanh, x=in_var, out=out_var, backward_func=tanh_grad_without_x,
    2.     skip_vars_in_backward_input=in_var)

    至此，使用 py_func 编写Python Op的步骤结束。我们可以与使用其他Op一样进行网路训练/预测。

    • py_func 的前向函数和反向函数内部不应调用 paddle.xx组网接口 ，因为前向函数和反向函数是在网络运行时调用的，而 paddle.xx 是在组建网络的阶段调用 。

    • skip_vars_in_backward_input 只能跳过前向输入变量和前向输出变量，不能跳过前向输出的梯度。

    • 若某个前向输出变量没有梯度，则 backward_func 将接收到 None 的输入。若某个前向输入变量没有梯度，则我们应在 backward_func 中主动返回 None。