自定义 Python 算子¶

动态图自定义 Python 算子¶

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

如何编写动态图 Python Op¶

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

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

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

forward和backward都是静态函数，它们的第一个参数是PyLayerContext对象。
backward 除了第一个参数以外，其他参数都是forward函数的输出Tensor的梯度，因此，backward输入的Tensor的数量必须等于forward输出Tensor的数量。如果您需在backward中使用forward中的Tensor，您可以利用save_for_backward和saved_tensor这两个方法传递Tensor。
backward的输出可以是Tensor或者list/tuple(Tensor)，这些Tensor是forward输入Tensor的梯度。因此，backward的输出Tensor的个数等于 forward 输入Tensor的个数。如果backward的某个返回值（梯度）在forward中对应的Tensor的stop_gradient属性为False，这个返回值必须是Tensor类型。

           import paddle
from paddle.autograd import PyLayer

# 通过创建`PyLayer`子类的方式实现动态图 Python Op
class cus_tanh(PyLayer):
    @staticmethod
    def forward(ctx, x):
        y = paddle.tanh(x)
        # ctx 为 PyLayerContext 对象，可以把 y 从 forward 传递到 backward。
        ctx.save_for_backward(y)
        return y

    @staticmethod
    # 因为 forward 只有一个输出，因此除了 ctx 外，backward 只有一个输入。
    def backward(ctx, dy):
        # ctx 为 PyLayerContext 对象，saved_tensor 获取在 forward 时暂存的 y。
        y, = ctx.saved_tensor()
        # 调用 Paddle API 自定义反向计算
        grad = dy * (1 - paddle.square(y))
        # forward 只有一个 Tensor 输入，因此，backward 只有一个输出。
        return grad

          

第二步：通过apply方法组建网络。 apply的输入为forward中除了第一个参数(ctx)以外的输入，apply的输出即为forward的输出。

           data = paddle.randn([2, 3], dtype="float32")
data.stop_gradient = False
# 通过 apply 运行这个 Python 算子
z = cus_tanh.apply(data)
z.mean().backward()

print(data.grad)

          

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

为了从forward到backward传递信息，您可以在forward中给PyLayerContext添加临时属性，在backward中读取这个属性。如果传递Tensor推荐使用save_for_backward和saved_tensor，如果传递非Tensor推荐使用添加临时属性的方式。

           import paddle
from paddle.autograd import PyLayer
import numpy as np

class tanh(PyLayer):
    @staticmethod
    def forward(ctx, x1, func1, func2=paddle.square):
        # 添加临时属性的方式传递 func2
        ctx.func = func2
        y1 = func1(x1)
        # 使用 save_for_backward 传递 y1
        ctx.save_for_backward(y1)
        return y1

    @staticmethod
    def backward(ctx, dy1):
        y1, = ctx.saved_tensor()
        # 获取 func2
        re1 = dy1 * (1 - ctx.func(y1))
        return re1

input1 = paddle.randn([2, 3]).astype("float64")
input2 = input1.detach().clone()
input1.stop_gradient = False
input2.stop_gradient = False
z = tanh.apply(x1=input1, func1=paddle.tanh)

          

forward 的输入和输出的类型任意，但是至少有一个输入和输出为Tensor类型。

           # 错误示例
class cus_tanh(PyLayer):
    @staticmethod
    def forward(ctx, x1, x2):
        y = x1+x2
        # y.shape: 列表类型，非 Tensor,输出至少包含一个 Tensor
        return y.shape

    @staticmethod
    def backward(ctx, dy):
        return dy, dy

data = paddle.randn([2, 3], dtype="float32")
data.stop_gradient = False
# 由于 forward 输出没有 Tensor 引发报错
z, y_shape = cus_tanh.apply(data, data)


# 正确示例
class cus_tanh(PyLayer):
    @staticmethod
    def forward(ctx, x1, x2):
        y = x1+x2
        # y.shape: 列表类型，非 Tensor
        return y, y.shape

    @staticmethod
    def backward(ctx, dy):
        # forward 两个 Tensor 输入，因此，backward 有两个输出。
        return dy, dy

data = paddle.randn([2, 3], dtype="float32")
data.stop_gradient = False
z, y_shape = cus_tanh.apply(data, data)
z.mean().backward()

print(data.grad)

          

如果 forward 的某个输入为Tensor且stop_gredient = True，则在backward中与其对应的返回值应为None。

           class cus_tanh(PyLayer):
    @staticmethod
    def forward(ctx, x1, x2):
        y = x1+x2
        return y

    @staticmethod
    def backward(ctx, dy):
        # x2.stop_gradient=True，其对应梯度需要返回 None
        return dy, None


data1 = paddle.randn([2, 3], dtype="float32")
data1.stop_gradient = False
data2 = paddle.randn([2, 3], dtype="float32")
z = cus_tanh.apply(data1, data2)
fake_loss = z.mean()
fake_loss.backward()
print(data1.grad)

          

如果 forward 的所有输入Tensor都是stop_gredient = True的，则backward不会被执行。

           class cus_tanh(PyLayer):
    @staticmethod
    def forward(ctx, x1, x2):
        y = x1+x2
        return y

    @staticmethod
    def backward(ctx, dy):
        return dy, None


data1 = paddle.randn([2, 3], dtype="float32")
data2 = paddle.randn([2, 3], dtype="float32")
z = cus_tanh.apply(data1, data2)
fake_loss = z.mean()
fake_loss.backward()
# 因为 data1.stop_gradient = True、data2.stop_gradient = True，所以 backward 不会被执行。
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 具体接口为：

           def py_func(func, x, out, backward_func=None, skip_vars_in_backward_input=None):
    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，则前向函数的定义格式为：

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

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

           import numpy as np

# 前向函数 1：模拟 tanh 激活函数
def tanh(x):
    # 可以直接将 Tensor 作为 np.tanh 的输入参数
    return np.tanh(x)

# 前向函数 2：将两个 2-D Tenosr 相加，输入多个 Tensor 以 list[Tensor]或 tuple(Tensor)形式
def element_wise_add(x, y):
    # 必须先手动将 Tensor 转换为 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):
    # 可以直接将 Tensor 作为 print 的输入参数
    print(x)

# 前向函数 1 对应的反向函数，默认的输入顺序为：x、out、out 的梯度
def tanh_grad(x, y, dy):
    # 必须先手动将 Tensor 转换为 numpy 数组，否则"+/-"等操作无法使用
    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 进行排除 :

           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

paddle.enable_static()

def create_tmp_var(program, name, dtype, shape):
    return program.current_block().create_var(name=name, dtype=dtype, shape=shape)

in_var = paddle.static.data(name='input', dtype='float32', shape=[-1, 28, 28])

# 手动创建前向输出变量
out_var = create_tmp_var(paddle.static.default_main_program(), name='output', dtype='float32', shape=[-1, 28, 28])

第三步：调用 py_func 组建网络

py_func 的调用方式为：

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

          

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

           paddle.static.nn.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 一样进行网路训练/预测。

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

py_func 的前向函数和反向函数内部不应调用 paddle.xx组网接口，因为前向函数和反向函数是在网络运行时调用的，而 paddle.xx 是在组建网络的阶段调用。
skip_vars_in_backward_input 只能跳过前向输入变量和前向输出变量，不能跳过前向输出的梯度。
若某个前向输出变量没有梯度，则 backward_func 将接收到 None 的输入。若某个前向输入变量没有梯度，则我们应在 backward_func 中主动返回 None。