如何写新的Python OP
PaddlePaddle Fluid通过 py_func
接口支持在Python端自定义OP。 py_func的设计原理在于Paddle中的LodTensor可以与numpy数组可以方便的互相转换,从而可以使用Python中的numpy API来自定义一个Python OP。
py_func接口概述
py_func
具体接口为:
- def py_func(func, x, out, backward_func=None, skip_vars_in_backward_input=None):
- pass
其中,
x
是Python Op的输入变量,可以是单个Variable
|tuple[Variable]
|list[Variable]
。多个Variable以tuple[Variable]或list[Variale]的形式传入,其中Variable为LoDTensor或Tenosr。out
是Python Op的输出变量,可以是单个Variable
|tuple[Variable]
|list[Variable]
。其中Variable既可以为LoDTensor或Tensor,也可以为numpy数组。func
是Python Op的前向函数。在运行网络前向时,框架会调用out = func(*x)
,根据前向输入x
和前向函数func
计算前向输出out
。在func
建议先主动将LoDTensor转换为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
中不需要的输入,可以是单个Variable
|tuple[Variable]
|list[Variable]
。
如何使用py_func编写Python Op
以下以tanh为例,介绍如何利用 py_func
编写Python Op。
- 第一步:定义前向函数和反向函数
前向函数和反向函数均由Python编写,可以方便地使用Python与numpy中的相关API来实现一个自定义的OP。
若前向函数的输入为 x_1
, x_2
, …, x_n
,输出为y_1
, y_2
, …, y_m
,则前向函数的定义格式为:
- def foward_func(x_1, x_2, ..., x_n):
- ...
- return y_1, y_2, ..., y_m
默认情况下,反向函数的输入参数顺序为:所有前向输入变量 + 所有前向输出变量 + 所有前向输出变量的梯度,因此对应的反向函数的定义格式为:
- def backward_func(x_1, x_2, ..., x_n, y_1, y_2, ..., y_m, dy_1, dy_2, ..., dy_m):
- ...
- return dx_1, dx_2, ..., dx_n
若反向函数不需要某些前向输入变量或前向输出变量,可设置 skip_vars_in_backward_input
进行排除(步骤三中会叙述具体的排除方法)。
注:,x_1, …, x_n为输入的多个LodTensor,请以tuple(Variable)或list[Variable]的形式在py_func中传入。建议先主动将LodTensor通过numpy.array转换为数组,否则Python与numpy中的某些操作可能无法兼容使用在LodTensor上。
此处我们利用numpy的相关API完成tanh的前向函数和反向函数编写。下面给出多个前向与反向函数定义的示例:
- import numpy as np
- # 前向函数1:模拟tanh激活函数
- def tanh(x):
- # 可以直接将LodTensor作为np.tanh的输入参数
- return np.tanh(x)
- # 前向函数2:将两个2-D LodTenosr相加,输入多个LodTensor以list[Variable]或tuple(Variable)形式
- def element_wise_add(x, y):
- # 必须先手动将LodTensor转换为numpy数组,否则无法支持numpy的shape操作
- x = np.array(x)
- y = np.array(y)
- if x.shape != y.shape:
- raise AssertionError("the shape of inputs must be the same!")
- result = np.zeros(x.shape, dtype='int32')
- for i in range(len(x)):
- for j in range(len(x[0])):
- result[i][j] = x[i][j] + y[i][j]
- return result
- # 前向函数3:可用于调试正在运行的网络(打印值)
- def debug_func(x):
- # 可以直接将LodTensor作为print的输入参数
- print(x)
- # 前向函数1对应的反向函数,默认的输入顺序为:x、out、out的梯度
- def tanh_grad(x, y, dy):
- # 必须先手动将LodTensor转换为numpy数组,否则"+/-"等操作无法使用
- return np.array(dy) * (1 - np.square(np.array(y)))
注意,前向函数和反向函数的输入均是 LoDTensor
类型,输出可以是Numpy Array或 LoDTensor
。 由于 LoDTensor
实现了Python的buffer protocol协议,因此即可通过 numpy.array
直接将 LoDTensor
转换为numpy Array来进行操作,也可直接将 LoDTensor
作为numpy函数的输入参数。但建议先主动转换为numpy Array,则可以任意的使用python与numpy中的所有操作(例如"numpy array的+/-/shape")。
tanh的反向函数不需要前向输入x,因此我们可定义一个不需要前向输入x的反向函数,并在后续通过 skip_vars_in_backward_input
进行排除 :
- def tanh_grad_without_x(y, dy):
- return np.array(dy) * (1 - np.square(np.array(y)))
- 第二步:创建前向输出变量
我们需调用 Program.current_block().create_var
创建前向输出变量。在创建前向输出变量时,必须指明变量的名称name、数据类型dtype和维度shape。
- import paddle.fluid as fluid
- def create_tmp_var(program, name, dtype, shape):
- return program.current_block().create_var(name=name, dtype=dtype, shape=shape)
- in_var = fluid.layers.data(name='input', dtype='float32', shape=[-1, 28, 28])
- # 手动创建前向输出变量
- out_var = create_tmp_var(fluid.default_main_program(), name='output', dtype='float32', shape=[-1, 28, 28])
- 第三步:调用
py_func
组建网络
py_func
的调用方式为:
- fluid.layers.py_func(func=tanh, x=in_var, out=out_var, backward_func=tanh_grad)
若我们不希望在反向函数输入参数中出现前向输入,则可使用 skip_vars_in_backward_input
进行排查,简化反向函数的参数列表。
- fluid.layers.py_func(func=tanh, x=in_var, out=out_var, backward_func=tanh_grad_without_x,
- skip_vars_in_backward_input=in_var)
至此,使用 py_func
编写Python Op的步骤结束。我们可以与使用其他Op一样进行网路训练/预测。
注意事项
py_func
的前向函数和反向函数内部不应调用fluid.layers.xxx
,因为前向函数和反向函数是在网络运行时调用的,且输入参数均为C++端的LoDTensor
; 而fluid.layers.xxx
是在组建网络的阶段调用的,且输入参数为Python端的Variable
。skip_vars_in_backward_input
只能跳过前向输入变量和前向输出变量,不能跳过前向输出的梯度。- 若某个前向输出变量没有梯度,则
backward_func
将接收到None
的输入。若某个前向输入变量没有梯度,则我们应在backward_func
中主动返回None
。