自定义Python算子

动态图自定义Python算子

Paddle 通过 PyLayer 接口和PyLayerContext接口支持动态图的Python端自定义OP。

相关接口概述

PyLayer 接口描述如下：

  1. class PyLayer:
  2.     @staticmethod
  3.     def forward(ctx, *args, **kwargs):
  4.         pass
  6.     @staticmethod
  7.     def backward(ctx, *args, **kwargs):
  8.         pass
  10.     @classmethod
  11.     def apply(cls, *args, **kwargs):
  12.         pass

其中，

  • forward 是自定义Op的前向函数，必须被子类重写，它的第一个参数是 PyLayerContext 对象，其他输入参数的类型和数量任意。

  • backward 是自定义Op的反向函数，必须被子类重写，其第一个参数为 PyLayerContext 对象，其他输入参数为forward输出Tensor的梯度。它的输出Tensor为forward输入Tensor的梯度。

  • apply 是自定义Op的执行方法，构建完自定义Op后，通过apply运行Op。

PyLayerContext 接口描述如下：

  1. class PyLayerContext:
  2.     def save_for_backward(self, *tensors):
  3.         pass
  5.     def saved_tensor(self):
  6.         pass

其中，

  • save_for_backward 用于暂存backward需要的Tensor，这个API只能被调用一次，且只能在forward中调用。

  • saved_tensor 获取被save_for_backward暂存的Tensor。

如何编写动态图Python Op

以下以tanh为例，介绍如何利用 PyLayer 编写Python Op。

  • 第一步：创建PyLayer子类并定义前向函数和反向函数

前向函数和反向函数均由Python编写，可以方便地使用Paddle相关API来实现一个自定义的OP。需要遵守以下规则：

  1. forward和backward都是静态函数，它们的第一个参数是PyLayerContext对象。

  2. backward 除了第一个参数以外，其他参数都是forward函数的输出Tensor的梯度，因此，backward输入的Tensor的数量必须等于forward输出Tensor的数量。如果您需在backward中使用forward中的Tensor，您可以利用save_for_backward和saved_tensor这两个方法传递Tensor。

  3. backward的输出可以是Tensor或者list/tuple(Tensor)，这些Tensor是forward输出Tensor的梯度。因此，backward的输出Tensor的个数等于forward输入Tensor的个数。如果backward的某个返回值（梯度）在forward中对应的Tensor是需要梯度，这个返回值必须是Tensor类型。

  1. import paddle
  2. from paddle.autograd import PyLayer
  4. # 通过创建`PyLayer`子类的方式实现动态图Python Op
  5. class cus_tanh(PyLayer):
  6.     @staticmethod
  7.     def forward(ctx, x):
  8.         y = paddle.tanh(x)
  9.         # ctx 为PyLayerContext对象，可以把y从forward传递到backward。
  10.         ctx.save_for_backward(y)
  11.         return y
  13.     @staticmethod
  14.     # 因为forward只有一个输出，因此除了ctx外，backward只有一个输入。
  15.     def backward(ctx, dy):
  16.         # ctx 为PyLayerContext对象，saved_tensor获取在forward时暂存的y。
  17.         y, = ctx.saved_tensor()
  18.         # 调用Paddle API自定义反向计算
  19.         grad = dy * (1 - paddle.square(y))
  20.         # forward只有一个Tensor输入，因此，backward只有一个输出。
  21.         return grad
  • 第二步：通过apply方法组建网络。 apply的输入为forward中除了第一个参数(ctx)以外的输入，apply的输出即为forward的输出。
  1. data = paddle.randn([2, 3], dtype="float32")
  2. data.stop_gradient = False
  3. # 通过 apply运行这个Python算子
  4. z = cus_tanh.apply(data)
  5. z.mean().backward()
  7. print(data.grad)

动态图自定义Python算子的注意事项

  • 为了从forward到backward传递信息，您可以在forward中给PyLayerContext添加临时属性，在backward中读取这个属性。如果传递Tensor推荐使用save_for_backward和saved_tensor，如果传递非Tensor推荐使用添加临时属性的方式。
  1. import paddle
  2. from paddle.autograd import PyLayer
  3. import numpy as np
  5. class tanh(PyLayer):
  6.     @staticmethod
  7.     def forward(ctx, x1, func1, func2=paddle.square):
  8.         # 添加临时属性的方式传递func2
  9.         ctx.func = func2
  10.         y1 = func1(x1)
  11.         # 使用save_for_backward传递y1
  12.         ctx.save_for_backward(y1)
  13.         return y1
  15.     @staticmethod
  16.     def backward(ctx, dy1):
  17.         y1, = ctx.saved_tensor()
  18.         # 获取func2
  19.         re1 = dy1 * (1 - ctx.func(y1))
  20.         return re1
  22. input1 = paddle.randn([2, 3]).astype("float64")
  23. input2 = input1.detach().clone()
  24. input1.stop_gradient = False
  25. input2.stop_gradient = False
  26. z = tanh.apply(x1=input1, func1=paddle.tanh)
  • forward的输入和输出的类型任意，但是至少有一个输入和输出为Tensor类型。
  1. # 错误示例
  2. class cus_tanh(PyLayer):
  3.     @staticmethod
  4.     def forward(ctx, x1, x2):
  5.         y = x1+x2
  6.         # y.shape: 列表类型，非Tensor,输出至少包含一个Tensor
  7.         return y.shape
  9.     @staticmethod
  10.     def backward(ctx, dy):
  11.         return dy, dy
  13. data = paddle.randn([2, 3], dtype="float32")
  14. data.stop_gradient = False
  15. # 由于forward输出没有Tensor引发报错
  16. z, y_shape = cus_tanh.apply(data, data)
  19. # 正确示例
  20. class cus_tanh(PyLayer):
  21.     @staticmethod
  22.     def forward(ctx, x1, x2):
  23.         y = x1+x2
  24.         # y.shape: 列表类型，非Tensor
  25.         return y, y.shape
  27.     @staticmethod
  28.     def backward(ctx, dy):
  29.         # forward两个Tensor输入，因此，backward有两个输出。
  30.         return dy, dy
  32. data = paddle.randn([2, 3], dtype="float32")
  33. data.stop_gradient = False
  34. z, y_shape = cus_tanh.apply(data, data)
  35. z.mean().backward()
  37. print(data.grad)
  • 如果forward的某个输入为Tensor且stop_gredient = True，则在backward中与其对应的返回值应为None。
  1. class cus_tanh(PyLayer):
  2.     @staticmethod
  3.     def forward(ctx, x1, x2):
  4.         y = x1+x2
  5.         return y
  7.     @staticmethod
  8.     def backward(ctx, dy):
  9.         # x2.stop_gradient=True，其对应梯度需要返回None
  10.         return dy, None
  13. data1 = paddle.randn([2, 3], dtype="float32")
  14. data1.stop_gradient = False
  15. data2 = paddle.randn([2, 3], dtype="float32")
  16. z = cus_tanh.apply(data1, data2)
  17. fake_loss = z.mean()
  18. fake_loss.backward()
  19. print(data1.grad)
  • 如果forward的所有输入Tensor都是stop_gredient = True的，则backward不会被执行。
  1. class cus_tanh(PyLayer):
  2.     @staticmethod
  3.     def forward(ctx, x1, x2):
  4.         y = x1+x2
  5.         return y
  7.     @staticmethod
  8.     def backward(ctx, dy):
  9.         return dy, None
  12. data1 = paddle.randn([2, 3], dtype="float32")
  13. data2 = paddle.randn([2, 3], dtype="float32")
  14. z = cus_tanh.apply(data1, data2)
  15. fake_loss = z.mean()
  16. fake_loss.backward()
  17. # 因为data1.stop_gradient = True、data2.stop_gradient = True，所以backward不会被执行。
  18. print(data1.grad is None)

静态图自定义Python算子

Paddle 通过 py_func 接口支持静态图的Python端自定义OP。 py_func的设计原理在于Paddle中的Tensor可以与numpy数组可以方便的互相转换，从而可以使用Python中的numpy API来自定义一个Python OP。

py_func接口概述

py_func 具体接口为：

  1. def py_func(func, x, out, backward_func=None, skip_vars_in_backward_input=None):
  2.     pass

其中，

  • 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] 。

如何使用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，则前向函数的定义格式为：

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

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

  1. def backward_func(x_1, x_2, ..., x_n, y_1, y_2, ..., y_m, dy_1, dy_2, ..., dy_m):
  2.     ...
  3.     return dx_1, dx_2, ..., dx_n

若反向函数不需要某些前向输入变量或前向输出变量，可设置 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. def element_wise_add(x, y):
  10.     # 必须先手动将Tensor转换为numpy数组，否则无法支持numpy的shape操作
  11.     x = np.array(x)  
  12.     y = np.array(y)
  14.     if x.shape != y.shape:
  15.         raise AssertionError("the shape of inputs must be the same!")
  17.     result = np.zeros(x.shape, dtype='int32')
  18.     for i in range(len(x)):
  19.         for j in range(len(x[0])):
  20.             result[i][j] = x[i][j] + y[i][j]
  22.     return result
  24. # 前向函数3：可用于调试正在运行的网络（打印值）
  25. def debug_func(x):
  26.     # 可以直接将Tensor作为print的输入参数
  27.     print(x)
  29. # 前向函数1对应的反向函数，默认的输入顺序为：x、out、out的梯度
  30. def tanh_grad(x, y, dy):
  31.     # 必须先手动将Tensor转换为numpy数组，否则"+/-"等操作无法使用
  32.     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”）。

tanh的反向函数不需要前向输入x，因此我们可定义一个不需要前向输入x的反向函数，并在后续通过 skip_vars_in_backward_input 进行排除 :

  1. def tanh_grad_without_x(y, dy):
  2.     return np.array(dy) * (1 - np.square(np.array(y)))
  • 第二步：创建前向输出变量

我们需调用 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 的调用方式为：

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

若我们不希望在反向函数输入参数中出现前向输入，则可使用 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一样进行网路训练/预测。

静态图自定义Python算子注意事项

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

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

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