【Feature Request】 PyLayer 功能支持动转静 @to_static 🚀

[Aurelius84]

需求描述 Feature Description

一、任务目标

请描述你正在做的项目是什么，如模型、论文、项目是什么？

期望飞桨的动态图中的 PyLayer 机制能够与飞桨的动转静@to_staitc 互通，支持模型中 PyLayer 的自定义层能够被 @to_static 感知并正确地生成静态图 Program，支撑转静训练和导出推理。

二、需求场景

请描述你的项目中为什么需要用此功能

飞桨的动态图 PyLayer 向用户提供了一种「高度灵活且便利」的自定义网络层前、反向计算的机制，比如存在一些用户自定义的计算逻辑，无法通过飞桨 现有的某个 API 或某些 API 组合实现，故可以借助 PyLayer 来实现用户的「idea」。如下是官网给出的一个样例：

import paddle
from paddle.autograd import PyLayer

# Inherit from PyLayer
class cus_tanh(PyLayer):
    @staticmethod
    def forward(ctx, x, func1, func2=paddle.square):
        # ctx is a context object that store some objects for backward.
        ctx.func = func2
        y = func1(x)
        # Pass tensors to backward.
        ctx.save_for_backward(y)
        return y

    @staticmethod
    # forward has only one output, so there is only one gradient in the input of backward.
    def backward(ctx, dy):
        # Get the tensors passed by forward.
        y, = ctx.saved_tensor()
        grad = dy * (1 - ctx.func(y))
        # forward has only one input, so only one gradient tensor is returned.
        return grad


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

print(data.grad)

目前，当动态图模型中包含 PyLayer 接口功能的使用时，暂不支持 @to_static 装饰以生成对应的静态图 Program，即 @to_static 暂时无法感知和处理 PyLayer 中 forward() 和 backward() 函数中的计算逻辑。

意味着，包含 PyLayer 的用户动态图模型，无法复用转静训练的「加速能力」，以及「导出推理」的功能。

三、功能描述

3.1 整体功能

飞桨的动态图中的 PyLayer 机制能够与飞桨的动转静@to_staitc 互通，支持模型中 PyLayer 的自定义层能够被 @to_static 感知并正确地生成静态图 Program，支撑转静训练和导出推理。

如下是一个功能支持后的全流程样例：

from symbol import parameters
import paddle
from paddle.autograd import PyLayer
from paddle.jit import to_static

# Inherit from PyLayer
class cus_tanh(PyLayer):
    @staticmethod
    def forward(ctx, x, func1, func2=paddle.square):
        # ctx is a context object that store some objects for backward.
        ctx.func = func2
        y = func1(x)
        # Pass tensors to backward.
        ctx.save_for_backward(y)
        return y

    @staticmethod
    # forward has only one output, so there is only one gradient in the input of backward.
    def backward(ctx, dy):
        # Get the tensors passed by forward.
        y, = ctx.saved_tensor()
        grad = dy * (1 - ctx.func(y))
        # forward has only one input, so only one gradient tensor is returned.
        return grad

class SimpleNet(paddle.nn.Layer):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.linear = paddle.nn.Linear(4, 8)
    
    @to_static
    def forward(self, x):
        y = self.linear(x)
        out = cus_tanh.apply(y, func1=paddle.tanh)
        out = paddle.mean(out)

        return out


def train(net, opt):
    batch_num = 10
    for i in range(batch_num):
        data = paddle.randn([2, 4])
        out = net(data)

        out.backward()
        opt.step()
        opt.clear_grad()
        print("loss: ", out.item())
    
    save(net)


def save(net):
    path = "simple_net"
    x_spec = paddle.static.InputSpec(shape=[-1, 4], dtype='float32', name='x')
    paddle.jit.save(net, path, input_spec=[x_spec])


if __name__ == '__main__':
    net = SimpleNet()
    sgd = paddle.optimizer.SGD(0.001, parameters=net.parameters())
    train(net, sgd)

PyLayer 在使用上需要遵循一定的规范，如：

• 子类必须包含静态的 forward 和 backward 函数，它们的第一个参数必须是 PyLayerContext
• 如果 backward 的某个返回值在 forward 中对应的 Tensor 是需要梯度，这个返回值必须为 Tensor
• backward 输入的 Tensor 的数量必须等于 forward 输出 Tensor 的数量
• 如果你需在 backward 中使用 forward 的输入 Tensor ，你可以将这些 Tensor 输入到 PyLayerContext 的 save_for_backward 方法，之后在 backward 中使用这些 Tensor 。
• backward 的输出 Tensor 的个数等于 forward 输入 Tensor 的个数

3.2 初期锚点

为了更好新增此 Feature，初期可以假设forward、backward 函数中都是基于Tensor的计算，且调用都是 Paddle API。，以此为「初期功能」的锚点，可以更专注于 Feature 的技术方案设计。

如下是初期锚点对应的一个简单样例：

from symbol import parameters
import paddle
from paddle.autograd import PyLayer
from paddle.jit import to_static

# Inherit from PyLayer
class cus_tanh(PyLayer):
    @staticmethod
    def forward(ctx, x): 
        # ctx is a context object that store some objects for backward.
        y = paddle.tanh(x)       # <------ 仅仅包含 Paddle API 的计算
        # Pass tensors to backward.
        ctx.save_for_backward(y)
        return y

    @staticmethod
    # forward has only one output, so there is only one gradient in the input of backward.
    def backward(ctx, dy):
        # Get the tensors passed by forward.
        y = ctx.saved_tensor()
        grad = dy * (1 - paddle.square(y))   # <------ 仅仅包含 Paddle API 的计算
        # forward has only one input, so only one gradient tensor is returned.
        return grad

class SimpleNet(paddle.nn.Layer):
    # .... （同上）

def train(net, opt):
    # .... （同上）

def save(net):
    # .... （同上）

if __name__ == '__main__':
    # .... （同上）

3.3 进阶考虑

在满足 「初期锚点」功能的技术方案设计后，可以再进阶考虑如果 forward 和 backward 函数中包含了其他更为高阶的用法，比如：

• 参数列表中不仅只包含了 Tensor，还包含了非 Tensor 对象
• forward 、 backward 函数调用了其他用户自定义的函数（这些函数里调用了Paddle API）

附：替代实现 Alternatives

PyLayer 功能本质上是为了支持用户「自定义」计算逻辑，解决飞桨现有 API 体系无法直接组合实现的问题（尤其是反向自定义）。

在飞桨框架中，存在一个等价的功能「自定义算子」是可以替换掉 PyLayer 的，且「自定义算子」功能是支持 动转静 @to_static 训练和导出推理的。

具体的用法可以参考飞桨官网：https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/guides/custom_op/index_cn.html

[Aurelius84]

如下是飞桨 PyLayer 和 动转静 @to_static 的技术实现原理。

一、PyLayer 实现原理

首先是 PyLayer 的机制流程：

PyLayer 继承自core.eager.PyLayer，提供了两个需要重写的接口：

class EagerPyLayer(
        with_mateclass(EagerPyLayerMeta, core.eager.PyLayer,
                       EagerPyLayerContext)):

    @staticmethod
    def forward(ctx, *args, **kwargs):
         raise NotImplementedError(
            "You must implement the forward function for PyLayer.")

    @staticmethod
    def backward(ctx, *args):
         raise NotImplementedError(
            "You must implement the backward function for PyLayer.")

C++ 端定义了一个 PyLayerObjec t的结构体：

typedef struct {
  PyObject_HEAD PyObject* container;
  bool container_be_packed;
  std::shared_ptr<egr::UnPackHookBase> unpack_hook;
  PyObject* non_differentiable;
  PyObject* not_inplace_tensors;
  bool materialize_grads;
  std::vector<bool> forward_input_tensor_is_duplicable;
  std::vector<bool> forward_output_tensor_is_duplicable;
  std::weak_ptr<egr::GradNodePyLayer> grad_node;
} PyLayerObject;

借助 BindEagerPyLayer() 函数实现了关键属性、方法的 bind：

void BindEagerPyLayer(PyObject* module) {
  auto heap_type = reinterpret_cast<PyHeapTypeObject*>(
      PyType_Type.tp_alloc(&PyType_Type, 0));
  heap_type->ht_name = ToPyObject("PyLayer");
  heap_type->ht_qualname = ToPyObject("PyLayer");
  auto type = &heap_type->ht_type;
  type->tp_name = "PyLayer";
  type->tp_basicsize = sizeof(PyLayerObject);
  type->tp_dealloc = (destructor)PyLayerDealloc;
  type->tp_methods = pylayer_methods;    // <----- 核心方法
  type->tp_getset = pylayer_properties;  // <----- 核心属性
  type->tp_new = (newfunc)PyLayerNew;
  // 省略
}

后端绑定的只有两个 method，一个是name()，另一个是 apply(self, *args, **kwargs) ，后者是 PyLayer 核心逻辑：

• 首先会解析用户传入的args和kwargs参数
• 然后调用用户的 forward函数：outputs = PyObject_Call(forward_fn, forward_args, kwargs) ;
• 创建反向grad_node节点：GradNodePyLayer
• 当用户调用loss.backward()时，最终会执行 GradNodePyLayer::operator()
  • 在C++端调用用户的backward函数：auto outputs = PyObject_CallObject(backward_fn, backward_args);

对于反向 GradNodePyLayer，其继承自GradNodeBase，核心成员如下：

class GradNodePyLayer : public GradNodeBase {
 public:
  GradNodePyLayer(PyObject* ctx,
                  size_t bwd_in_slot_num,
                  size_t bwd_out_slot_num)
      : GradNodeBase(bwd_in_slot_num, bwd_out_slot_num) {
    ctx_ = ctx;
    Py_INCREF(ctx_);
  }
  
  private:
  PyObject* ctx_{nullptr};    //<----- 记录了Pythond端的PyLayer对象，用于获取backward函数指针
  std::vector<std::vector<phi::DenseTensorMeta>> forward_outputs_meta_;
  std::vector<std::vector<paddle::platform::Place>> forward_outputs_place_;
};

[Aurelius84]

二、动转静 实现原理

具体可以参考官方文档：https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/guides/jit/principle_cn.html

在飞桨框架内部，动转静模块在转换上主要包括对输入数据的 InputSpec 的处理，对函数调用的递归转写，对 IfElse、For、While 控制语句的转写，以及 Layer 的 Parameters 和 Buffers 变量的转换。下面将介绍动转静模块的转换过程。

上图是动转静转换和训练执行的基本流程：

2.1 AST 解析动态图代码

当某个函数被 @to_static 装饰、或用 paddle.jit.to_static() 包裹时，飞桨会隐式地解析动态图的 Python 代码（即解析：抽象语法树，简称 AST）。

2.2 AST 转写，得到静态图代码

• 函数转写：递归地对所有函数进行转写，实现用户仅需在最外层函数添加 @to_static 的体验效果。

• 控制流转写：用户的代码中可能包含依赖 Tensor 的控制流代码，飞桨框架会自动且有选择性地将 if、for、while 转换为静态图对应的控制流。

• 其他语法处理：包括 break、continue、assert、提前 return 等语法的处理。

2.3 生成静态图的 Program 和 Parameters

得到静态图代码后，根据用户指定的 InputSpec 信息（或训练时根据实际输入 Tensor 隐式创建的 InputSpec ）作为输入，执行静态图代码生成 Program。每个被装饰的函数，都会被替换为一个 StaticFunction 对象，其持有此函数对应的计算图 Program，在执行 paddle.jit.save 时会被用到。

对于trainable=True 的 Buffers 变量，动转静会自动识别并将其和 Parameters 一起保存到 .pdiparams 文件中。

2.4 执行动转静训练

使用执行引擎执行函数对应的 Program，返回输出 out。

执行时会根据用户指定的 build_strategy 策略应用图优化技术，提升执行效率。

2.5 使用 paddle.jit.save 保存静态图模型

使用 paddle.jit.save 时会遍历模型 net 中所有的函数，将每个的 StaticFunction 中的计算图 Program 和涉及到的 Parameters 序列化为磁盘文件。

[Aurelius84]

一些技术思路

目前框架中动态图中 PyLayer、静态图PyFunc OP 机制现状如下表：

如果要支持 PyLayer 动转静机制，需要框架端新增哪些组件或功能，难度如何？

1. Op静态化。

即是否必须要有静态图的前向、反向OpDesc，Kernel 的对应实现 ？

待思考点：

• 用户定义的函数是基于 Tensor的，新增 Op 其 kernel内 部需要做Variable -> Tensor的来回转换

2. 函数序列化。

即是否必须要支持用户自定义的函数能序列化、反序列化，打通Inference、或Lite ？

待思考点：

• 用户的函数非常多样性，是否可以做到完备序列化和反序列化。
• 训练端支持后，inference 和 Lite 扩展支持的技术性风险依旧较大。

[Fancy-hjyp]

关于如何使 PyLayer 功能支持动转静 @to_static 的一些想法

1.如何让@to_static感知到PyLayer？

@to_static本质上是通过Python的装饰器将一个函数转写为一个StaticFunction对象，该装饰器主要工作内容为：通过Python AST转写为静态图代码，将动态图代码转写为静态图代码，例如将ast.call，ast.if类型的结点替换为静态图下对应的语法代码。

为了让@to_static感知到PyLayer，可以新定义一个装饰器，比如叫@to_static_pylayer，在新定义的装饰器中使用to_static()函数将动态图代码转换为静态图代码。

还有一种方法是直接修改to_static()的内部实现，在内部实现时判断调用的函数是否使用了xxx.apply()这种方式来调用，因为对于继承自PyLayer的自定义python算子，在使用时必须通过xxx.apply()方式调用。

通过上面的步骤，我们可以完成动态代码到静态代码的转写。

下面是一个直接使用@to_static转写的例子。

import paddle
from paddle.autograd import PyLayer
from paddle.jit import to_static

class cus_tanh(PyLayer):
    @staticmethod
    @to_static
    def forward(ctx, x):
        y = paddle.tanh(x)
        ctx.save_for_backward(y)
        return y

    @staticmethod
    @to_static
    def backward(ctx, dy):
        y, = ctx.saved_tensor()
        grad = dy * (1 - paddle.square(y))
        return grad

model = cus_tanh()
result = model.apply(paddle.to_tensor([1., 2., 3.]))
print(cus_tanh.forward.code)
print(cus_tanh.backward.code)
result.mean().backward()

# output
def forward(ctx, x):
    __return_value_0 = None
    y = paddle.tanh(x)
    _jst.Call(ctx.save_for_backward)(y)
    __return_value_0 = _jst.Ld(y)
    return _jst.Ld(__return_value_0)

def backward(ctx, dy):
    __return_value_1 = None
    y, = _jst.Call(ctx.saved_tensor)()
    grad = _jst.Ld(dy) * (1 - paddle.square(y))
    __return_value_1 = _jst.Ld(grad)
    return _jst.Ld(__return_value_1)

2.转写时需要注意的问题？

因为静态图的输入必须是Variable变量，不能是Tensor，静态图下Paddle执行Block.append_op()函数时首先会使用param_guard()函数将Tensor转为Variable，然后在执行静态图创建逻辑。

所以第1步转换的代码也需要在开始的时候将Tensor转为Variable。

3.静态图的导出

最后还需要考虑将创建好的静态图进行导出，所以要支持持久化功能，以便用来训练和推理。

有几个问题还不太清楚：

问题1：静态图创建流程是如何的？

我的理解是先执行一遍python端转写后静态图代码，执行的过程就相当于在c++端创建好了整个静态图，然后是整个静态图的优化加速，最后是在整张图上的训练推理。官方文档上说静态图的优势在于减少了python与c++之间的交互，实际上是指创建好整张图后，直接在c++端进行了运算，不知道这样的理解对不对？

基于上面的理解，PyLayer是在Python端定义的算子。动态图模式下，因为不会预先建立整张图，所以Paddle自带的算子是在C++端执行的，运行用户自定义的Python算子还是会回到Python端执行，存在一个执行端切换的过程。

问题2：转写后的PyLayer代码会直接执行吗？

假设我们得到了转写后静态图对应的PyLayer代码，还是像之前一样，先执行一遍Python端代码来在c++端创建整张静态图。如果执行到转写后的PyLayer代码，也应该会创建对应的静态图结点吧？如果创建的话，整个静态图的执行过程就是完整的。

问题3：对于Paddle不支持的运算是指什么呢？

所有的运算不都是建立在加减乘除、卷积、池化这些基础操作上吗，对于issue里面初期可以假设forward、backward 函数中都是基于Tensor的计算，且调用都是 Paddle API与基本运算的区别不太理解，是指执行非Paddle API不会在c++端建立相应的结点吗，只是在Python端进行了运算？

[Aurelius84]

@Tomoko-hjf

问题1：静态图创建流程是如何的？

答：静态图与动态图是共用的一套API体系，每个API下都有两个分支：动态图和静态图。静态图下在逐个执行python API 时，是不会真正执行的，只是依次记录每个Op，来生成完整的前向的Program，真正执行是通过动转静转交给底层执行器来执行的。

问题2：转写后的PyLayer代码会直接执行吗？
答：基于此，如果要考虑支持PyLayer：①静态图模式下执行到 pylayer.apply()时，此时需要感知forward 里的组网代码并依次执行记录Op到外层的Program中 ②同时标记backward函数逻辑（类似Hook？），在调用AutoDiff生成反向时到达PyLayer处，将此backward里的Op 放到反向逻辑里。

以上仅是我目前有限视野下的思路，欢迎一起讨论