PR types

New features

PR changes

Others

Description

Original requirements

It is expected that the PyLayer mechanism in the dynamic graph of the Paddle can interoperate with the @to_staitc of the Paddle's dynamic to static, supporting the custom layer of the PyLayer in the model to be perceived by the @to_static and correctly generating the static graph Program, and supporting the training of the dynamic to static and the exporting for inference.

So our main objective is to Enable PyLayer to Support @to_static. In this PR, we support @to_static to percept and translate the logic of forward function and backward function from PyLayer and also make effort to percept and process PyLayerContext in dy2st module.

Related issues

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

Related PR

[GLCC]Part-1: Add pylayer op to Support @to_static: #56108

Project Records

https://github.com/MarioLulab/GLCC2023-Paddle-Record

Main works

1. Enable @to_static to sense and handle the computational logic in the forward() and backward() functions in dygraph PyLayer, by modifying the function convert_load and returning StaticPyLayer. StaticPyLayer can complete dy2st code translation, and StaticPyLayer.apply can create forward block and backward block.
2. We subtitude PyLayerContext with StaticPyLayerContext to percept the logic of PyLayerContext. And we use the STEP_SCOPES and skip_gc_vars to solve the problem of pass intermediate variables between forward and backward. We modify the code of garbage collection to support keeping neccessary Variables in step_scope after finishing executing forward graph.
3. Add some testcases of @to_static support for PyLayer and static_pylayer API
4. Complete some TODO work and polish the code mentioned by the comments of [GLCC]Part-1: Add pylayer op to Support @to_static #56108

Future works

1. Refine the prune.cc to support pylayer op for jit.save() for inference.

2. static_pylayer API support not only positional argument but also keyword argument.

3. static_pylayer API support not only paddle.Tensor but also non-Tensor objects will be included in the argument list.

4. static_pylayer API support stop_gradient = False for some inputs and stop_gradient = True for the other inputs.

5. StaticPyLayerContext support more features (e.g. mark_not_inplace) which align with dygraph PyLayer