Customizable Python Layer in Dygraph #32130

hbwx24 · 2021-04-07T11:52:51Z

PR types

Function optimization

PR changes

APIs

Describe

功能描述：
为了支持动态图Python端自定义OP的功能，添加了如下API

class PyLayer: 通过继承PyLayer构造一个子类的方式实现自定义的Layer。这个子类需要遵循以下规则：
- 子类必须包含forward和backward函数，backward和forward必须用@staticmethod修饰。它们的第一个参数必须是context，并且它们的返回结果中不能包含None。
- backward第一个输入为context，其他输入为forward输出的tensor的梯度，因此backward输入tensor个数 = forward输出tensor个数。若需要forward的输入输出计算梯度，可以使用ctx.save_for_backward存储需要的tensor，之后可以backward中使用。
- backward输出只能为Tensor 或者tuple/list[Tensor]，因为backward的输出tensor是forward输入tensor的梯度，因此backward输出tensor个数 = forward输入tensor个数。
- 使用apply调用自定义的layer。
PyLayerContext:可以利用它把某些对象从forward传递到backward。具体用法在下面示例代码中。
- PyLayerContext.save_for_backward:在forward使用，将tensor存储到PyLayerContext对象中。
- PyLayerContext.saved_tensor：从PyLayerContext对象中获取保存的tensor。
上述API添加在paddle.autograd下，通过继承paddle.autograd.PyLayer构造自定义的Layer。

TODO：

目前只支持最基础的功能，还有待进一步完善。
目前只支持动态图，后续根据需求添加与静态图相关的功能。

示例：

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


class double_tanh(PyLayer):
    @staticmethod
    def forward(ctx, x1, x2, func1, func2=paddle.square):
          ctx.func = func2
          y1 = func1(x1)
          y2 = func1(x2)
          # 保存backward需要的tensor
          ctx.save_for_backward(y1, y2)
          return y1, y2

    @staticmethod
    def backward(ctx, dy1, dy2):
        # dy1：y1的梯度；dy2：y2的梯度。
        # 获取在forward中保存的tensor
        y1, y2 = ctx.saved_tensor()
        re1 = dy1 * (1 - ctx.func(y1))
        re2 = dy2 * (1 - paddle.square(y2))
        return re1, re2


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

# 使用自定义layer
z = double_tanh.apply(input1, input1, paddle.tanh, paddle.square)

z = z[0] + z[1]
z.mean().backward()

z2 = paddle.tanh(input2) + paddle.tanh(input2)
z2.mean().backward()
print(np.array_equal(input1.grad,input2.grad))
print(input1.grad-input2.grad)

API的英文文档：
http://10.136.157.23:8090/documentation/docs/en/api/paddle/autograd/py_layer/PyLayer_en.html

http://10.136.157.23:8090/documentation/docs/en/api/paddle/autograd/py_layer/PyLayerContext_en.html

chenwhql · 2021-04-08T06:51:07Z

paddle/fluid/imperative/dygraph_grad_maker.h

@@ -279,6 +279,8 @@ class TracedGradOp {

  void SetType(const std::string& type) { op_->SetType(type); }

+  const framework::OperatorBase& InnerOp() { return op_->InnerOp(); }


add const limit after InnerOp() ?

chenwhql · 2021-04-08T06:53:30Z

paddle/fluid/imperative/layer.h

@@ -279,6 +279,6 @@ std::shared_ptr<GradOpNode> CreateGradOpNode(
    const NameVarBaseMap& outs, const framework::AttributeMap& attrs,
    const platform::Place& place,
    const std::map<std::string, std::string>& inplace_map);
-
+void ClearNoNeedBufferInputs(OpBase* op);


add blank line before and after this statement?

chenwhql · 2021-04-08T07:05:44Z

paddle/fluid/imperative/py_layer_fwd/py_layer_fwd.h

+          auto a = ptr->cast<std::shared_ptr<VarBase>>();
+          input_vars.push_back(a);
+        }
+      } catch (py::cast_error& err) {


这里需要报错吧，可以提示现在不支持，但跳过是不是不太好，会让用户困惑？

添加了注释，收集所有的tensor变量，忽略所有非tensor变量。

chenwhql · 2021-04-08T07:05:51Z

paddle/fluid/imperative/py_layer_fwd/py_layer_fwd.h

+          auto a = ptr->second.cast<std::shared_ptr<VarBase>>();
+          input_vars.push_back(a);
+        }
+      } catch (py::cast_error&) {


chenwhql · 2021-04-08T07:07:08Z

paddle/fluid/imperative/py_layer_fwd/py_layer_fwd.h

+          output_vars.push_back(temp_out);
+        } catch (py::cast_error&) {
+          PADDLE_THROW(platform::errors::Unimplemented(
+              "The output of forward should be `Tensor`."));


PyLayer.forward？写全一点？让用户清楚是哪个forword，下面有多处相同报错，建议都完善一下

chenwhql · 2021-04-08T07:37:46Z

paddle/fluid/operators/py_layer_op.h

+class PyLayerOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+  void InferShape(framework::InferShapeContext* ctx) const override {}


如果不支持的话，加PADDLE_THROW，不要为空，否则误用的时候不太容易找问题，可以加几个异常单测，覆盖一下

添加Log提醒这里没有实现Infershape。

chenwhql · 2021-04-08T07:42:14Z

paddle/fluid/operators/py_layer_op.h

+ protected:
+  framework::OpKernelType GetExpectedKernelType(
+      const framework::ExecutionContext& ctx) const override {
+    auto data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X");


这里预期就是不支持不同类型运算的吧，这里X是Duplicable的，如果X的两个输入类型不一致，比如一个是float，一个是int，会报错吧

是的，不支持不同类型的运算，如果两个类型不一致会触发警告：『The dtype of left and right variables are not the same, left dtype is paddle.int64, but right dtype is paddle.float32, the right dtype will convert to paddle.int64』

chenwhql · 2021-04-08T07:42:40Z

paddle/fluid/operators/py_layer_op.cc

+                  ops::PyLayerGradOpMaker<paddle::framework::OpDesc>);
+
+REGISTER_OP_CPU_KERNEL(
+    py_layer, ops::PyLayerOpKernel<paddle::platform::CPUDeviceContext, float>,


数据类型需要补全吗？这个要支持所有类型吧

chenwhql · 2021-04-08T07:44:23Z

python/paddle/nn/layer/py_layer.py

+        return super(FunctionMeta, cls).__init__(name, bases, attrs)
+
+
+class PyLayer(with_mateclass(FunctionMeta, CFunction)):


先移到autograd下吧，补全文档

chenwhql · 2021-04-08T07:47:48Z

paddle/fluid/imperative/py_layer_fwd/py_layer_fwd.h

@@ -0,0 +1,158 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.


这里也是，新建目录只管理了一个文件，建议直接移到imperative下，不需要加目录

chenwhql · 2021-04-09T03:18:34Z

python/paddle/autograd/py_layer.py

+    """
+    Build a custom `Layer` by creating subclasses. Subclasses need to follow the following rules:
+        1. Subclasses contain `forward` and `backward` function. Both forward and backward are @staticmethod.
+            Their first argument should be a context and `None` can be included in the returned result.


None can还是can not？

thx，done.

chenwhql · 2021-04-09T03:20:12Z

python/paddle/autograd/py_layer.py

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


custom layer -> PyLayer?

chenwhql · 2021-04-09T03:20:20Z

python/paddle/autograd/py_layer.py

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


TCChenlong

LGTM

chenwhql

LGTM

…or_bk

XiaoguangHu01

LGTM

TCChenlong

LGTM

ForFishes

LGTM

TCChenlong

LGTM

hbwx24 added 13 commits April 3, 2021 05:22

custom python backward

e48e0d1

polish up the code

fb2e34b

polish up the code

6fbcf4e

polish up the code.

b56fda6

Merge branch 'develop' into custom-op/cus_backward

e73ca66

Fix code format and comments.

0c594bd

Delete redundant files.

baf4793

add unnittest.

4f52f6c

edit unnittest.

14d6080

edit unnittest.

7ac0403

Remove redundant header files.

3e5ae14

Improve coverage and remove redundant code.

01557fc

support saving for backward.

dd592db

chenwhql reviewed Apr 8, 2021

View reviewed changes

hbwx24 added 2 commits April 8, 2021 13:21

polish code according to comments.

a774fa9

Add support type for PyLayer.

4a69bdc

chenwhql reviewed Apr 9, 2021

View reviewed changes

hbwx24 added 2 commits April 9, 2021 04:39

Modify the DOC.

d6b7f24

Merge branch 'develop' into custom-op/save_for_bk

72b5ab1

hbwx24 changed the title ~~Custom op/save for bk~~ python 端自定义OP Apr 9, 2021

hbwx24 added 6 commits April 9, 2021 06:26

polish Doc.

1b1a5cc

polish Doc.

19fe271

polish Doc.

3e27b57

polish Doc.

24a78e2

polish Doc.

3cf1cc0

polish Doc.

3813c3d

TCChenlong previously approved these changes Apr 9, 2021

View reviewed changes

hbwx24 changed the title ~~python 端自定义OP~~ Customizable Python Layer Apr 9, 2021

hbwx24 changed the title ~~Customizable Python Layer~~ Customizable Python Layer in Dygraph Apr 9, 2021

chenwhql previously approved these changes Apr 9, 2021

View reviewed changes

lanxianghit previously approved these changes Apr 13, 2021

View reviewed changes

XieYunshen previously approved these changes Apr 13, 2021

View reviewed changes

raindrops2sea previously approved these changes Apr 13, 2021

View reviewed changes

hbwx24 added 2 commits April 14, 2021 06:41

polish code and make the code robust.

e219513

Merge remote-tracking branch 'upstream/develop' into custom-op/save_f…

c7405d3

…or_bk

hbwx24 dismissed stale reviews from raindrops2sea, XieYunshen, lanxianghit, ForFishes, chenwhql, and TCChenlong via c7405d3 April 14, 2021 06:47

XiaoguangHu01 previously approved these changes Apr 14, 2021

View reviewed changes

XieYunshen previously approved these changes Apr 14, 2021

View reviewed changes

lanxianghit previously approved these changes Apr 14, 2021

View reviewed changes

TCChenlong previously approved these changes Apr 14, 2021

View reviewed changes

ForFishes previously approved these changes Apr 14, 2021

View reviewed changes

Modify the code format.

f01faec

hbwx24 dismissed stale reviews from ForFishes, TCChenlong, lanxianghit, XieYunshen, and XiaoguangHu01 via f01faec April 14, 2021 08:27

TCChenlong approved these changes Apr 14, 2021

View reviewed changes

lanxianghit approved these changes Apr 14, 2021

View reviewed changes

XieYunshen approved these changes Apr 14, 2021

View reviewed changes

raindrops2sea approved these changes Apr 14, 2021

View reviewed changes

chenwhql approved these changes Apr 15, 2021

View reviewed changes

chenwhql merged commit 29f6522 into PaddlePaddle:develop Apr 15, 2021

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Customizable Python Layer in Dygraph #32130

Customizable Python Layer in Dygraph #32130

hbwx24 commented Apr 7, 2021 •

edited

Loading

chenwhql Apr 8, 2021

hbwx24 Apr 8, 2021

chenwhql Apr 8, 2021

hbwx24 Apr 8, 2021

chenwhql Apr 8, 2021

hbwx24 Apr 8, 2021

chenwhql Apr 8, 2021

hbwx24 Apr 8, 2021

chenwhql Apr 8, 2021

hbwx24 Apr 8, 2021

chenwhql Apr 8, 2021

hbwx24 Apr 8, 2021

chenwhql Apr 8, 2021

hbwx24 Apr 8, 2021

chenwhql Apr 8, 2021

hbwx24 Apr 8, 2021 •

edited

Loading

chenwhql Apr 8, 2021

hbwx24 Apr 8, 2021

chenwhql Apr 8, 2021

hbwx24 Apr 8, 2021

chenwhql Apr 9, 2021

hbwx24 Apr 9, 2021

chenwhql Apr 9, 2021

hbwx24 Apr 9, 2021

chenwhql Apr 9, 2021

hbwx24 Apr 9, 2021

TCChenlong left a comment

chenwhql left a comment

XiaoguangHu01 left a comment

TCChenlong left a comment

ForFishes left a comment

TCChenlong left a comment

		@@ -279,6 +279,8 @@ class TracedGradOp {

		void SetType(const std::string& type) { op_->SetType(type); }

		const framework::OperatorBase& InnerOp() { return op_->InnerOp(); }

		return super(FunctionMeta, cls).__init__(name, bases, attrs)


		class PyLayer(with_mateclass(FunctionMeta, CFunction)):

		@@ -0,0 +1,158 @@
		// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.

Customizable Python Layer in Dygraph #32130

Customizable Python Layer in Dygraph #32130

Conversation

hbwx24 commented Apr 7, 2021 • edited Loading

PR types

PR changes

Describe

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hbwx24 Apr 8, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

TCChenlong left a comment

Choose a reason for hiding this comment

chenwhql left a comment

Choose a reason for hiding this comment

XiaoguangHu01 left a comment

Choose a reason for hiding this comment

TCChenlong left a comment

Choose a reason for hiding this comment

ForFishes left a comment

Choose a reason for hiding this comment

TCChenlong left a comment

Choose a reason for hiding this comment

hbwx24 commented Apr 7, 2021 •

edited

Loading

hbwx24 Apr 8, 2021 •

edited

Loading