Refactor and simplify hook design & add Tensor.register_hook API #31775

chenwhql · 2021-03-22T05:15:06Z

PR types

New features

PR changes

APIs

Describe

Refactor and simplify hook design & add Tensor.register_hook API

1. Refactor

Simplify Hook class design

original classes

- OpBasePreHook
  - PyOpBasePreHook (Implement later)
  - CppOpBasePreHook (Implement later)
- GradAccumulatorPostHook
  - PyGradAccumulatorPostHook (Implement later)
  - CppGradAccumulatorPostHook (Implement later)
  - LambdaGradAccumulatorPostHook
- InteriorVarHookPipeline
- LeafVarHookPipeline

new classes

- VariableWrapperHook
  - PyVariableWrapperHook
- InplaceVariableWrapperHook
  - PyInplaceVariableWrapperHook (Implement later)
  - LambdaInplaceVariableWrapperHook

The input of hook operator is VariableWrapper, so hook is completely managed by VariableWrapper itself
Remove weak_ptr in OpBase and GradientAccumullator
Remove several hook related methods

2. Add Tensor.register_hook method

Support register backward hook in Python

from __future__ import print_function

import paddle

# hook function return None
def print_hook_fn(grad):
    print(grad)

# hook function return Tensor
def double_hook_fn(grad):
    grad = grad * 2
    return grad

x = paddle.to_tensor([0., 1., 2., 3.], stop_gradient=False)
y = paddle.to_tensor([4., 5., 6., 7.], stop_gradient=False)
z = paddle.to_tensor([1., 2., 3., 4.])

# one Tensor can register multiple hooks
h = x.register_hook(print_hook_fn)
x.register_hook(double_hook_fn)

w = x + y
# register hook by lambda function
w.register_hook(lambda grad: grad * 2)

o = z.matmul(w)
o.backward()
# print_hook_fn print content in backward
# Tensor(shape=[4], dtype=float32, place=CUDAPlace(0), stop_gradient=False,
#        [2., 4., 6., 8.])

print("w.grad:", w.grad) # w.grad: [1. 2. 3. 4.] - no changed
print("x.grad:", x.grad) # x.grad: [ 4.  8. 12. 16.] - deal with by two *2 hook
print("y.grad:", y.grad) # y.grad: [2. 4. 6. 8.] - - deal with by one *2 hook

# remove hook
h.remove()

3. Doc

related cn doc: PaddlePaddle/docs#3390

英文由于文档抽取问题，现在无法预览

paddle-bot-old · 2021-03-22T05:15:08Z

Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

JiabinYang

some comments

JiabinYang · 2021-03-29T07:00:27Z

paddle/fluid/imperative/basic_engine.cc

@@ -408,9 +412,25 @@ void BasicEngine::Execute() {
        }
      }

+      for (auto& pair : tmp_ins) {


how about create tmp_ins only when it needed, it seems make too many tmp variable_wrapper copy here

JiabinYang · 2021-03-29T07:03:03Z

paddle/fluid/imperative/basic_engine.cc

-          accumulator->CallBackwardPostHooks();
-        }
+        // 3. Call backward Hooks for `var_`
+        accumulator->CallReduceHooks();


Bad name, or may use inherent to fix it? CallHooks indicates invoke all hooks, but CallReduceHooks make it confused to me

done, CallHooks-> CallGradientHooks

JiabinYang · 2021-03-29T07:45:07Z

paddle/fluid/imperative/gradient_accumulator.cc

+                    platform::errors::InvalidArgument("Leaf Tensor's inner var "
+                                                      "is not initialized when "
+                                                      "call gradient hook."));
+  if (var_->HasHook()) {


seal this or make it has difference with the same code in Execute

only for loop is similar

JiabinYang · 2021-03-29T07:45:36Z

paddle/fluid/imperative/gradient_accumulator.cc

+  }
+}
+
+void GradientAccumulator::CallReduceHooks() {


do some check to differ it with normal hook

JiabinYang · 2021-03-29T07:47:24Z

paddle/fluid/imperative/gradient_accumulator.h

+   *    parallel multi-card training.
+   */
+
+  void CallHooks();


make this two func not a parallel structure with related name

JiabinYang · 2021-03-29T07:48:35Z

paddle/fluid/imperative/hooks.h

 */
-class OpBasePreHook {
+class VariableWrapperHook {


how about make a abstract class of Hook to seal different kinds of hooks

I have tried, that is not a good idea

JiabinYang · 2021-03-29T07:51:10Z

paddle/fluid/imperative/variable_wrapper.h

+  int64_t next_hook_id_{0};
+  // Hooks used to register hook for grad var, support adding and removing,
+  // key is the accumulated int64_t value
+  std::map<int64_t, std::shared_ptr<VariableWrapperHook>> hooks_;


why map here

the hook remove helper need to hold hook id for removing it correctlly

… hook/refactor_hook_impl_and_add_py_api

ForFishes

LGTM

zhwesky2010 · 2021-03-30T08:48:58Z

paddle/fluid/imperative/gradient_accumulator.cc

-   * If the gradient has been calculated by previous graph,
-   * it should be added to the previous graph result.
+   * If the leaf gradient has been calculated done, the inner_var_
+   * should be added to the var_.
   */
  if (!var_->IsLeafGrad() || !SumGradCompleted() || !HasInnerVar()) {


!HasInnerVar() 这个应该能去掉了

这个不能吧，现在每次调用AccumulatedGrad仍然要求有InnerVar的

zhwesky2010 · 2021-03-30T08:56:06Z

paddle/fluid/imperative/gradient_accumulator.cc

+    for (const auto& hook_pair : var_->GetHooks()) {
+      tmp_var = (*hook_pair.second)(tmp_var);
+    }
+    inner_var_ = tmp_var;


叶子节点在GradientAccumulator里面做CallGradientHooks就会替代自己内部的inner_var_，相当于inplace了吧

是的，本来就是inplace的，这里改成这样，主要目的是统一hook的基类管理和调用，如果这里使用InplaceHook，那之前的HookPipeLine那些就仍然需要，数据结构和逻辑都会比较复杂

TCChenlong

LGTM

JiabinYang

LGTM

… hook/refactor_hook_impl_and_add_py_api

refactor and simplify hook design

c0b947f

chenwhql added 10 commits March 23, 2021 02:20

fix reducer add hook error

b4b3e9f

add Tensor.register_hook basic impl

16b3dcd

refine prepare data impl

2553179

revert prepare data change

2fac74f

support register_hook for Tensor

de8b2df

add hook test in model

665b15b

polish tests and doc example

118cc07

fix double grad test failed

aa68578

remove reduce hook func

e8f799a

fix set empty error

21eceec

chenwhql mentioned this pull request Mar 29, 2021

Add Tensor.register_hook cn doc PaddlePaddle/docs#3390

Merged

chenwhql requested review from JiabinYang, phlrain, zhwesky2010 and ForFishes March 29, 2021 06:39

JiabinYang reviewed Mar 29, 2021

View reviewed changes

chenwhql added 4 commits March 29, 2021 09:25

polish code by comments

d5468e5

change reduce_hook to mutable_hook

c0838dc

remove useless tmp_ins

dbd3c34

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

fe79a89

… hook/refactor_hook_impl_and_add_py_api

ForFishes previously approved these changes Mar 30, 2021

View reviewed changes

zhwesky2010 reviewed Mar 30, 2021

View reviewed changes

zhwesky2010 previously approved these changes Mar 30, 2021

View reviewed changes

TCChenlong previously approved these changes Mar 31, 2021

View reviewed changes

JiabinYang previously approved these changes Mar 31, 2021

View reviewed changes

chenwhql requested a review from lanxianghit March 31, 2021 03:16

lanxianghit previously approved these changes Mar 31, 2021

View reviewed changes

fix shape code format error

7c9fd70

chenwhql dismissed stale reviews from lanxianghit, JiabinYang, TCChenlong, zhwesky2010, and ForFishes via 7c9fd70 March 31, 2021 11:34

chenwhql added 2 commits March 31, 2021 11:50

Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into…

11c26a9

… hook/refactor_hook_impl_and_add_py_api

fix shape code format error

ef087a5

TCChenlong approved these changes Mar 31, 2021

View reviewed changes

lanxianghit approved these changes Mar 31, 2021

View reviewed changes

chenwhql merged commit dbeb3ea into PaddlePaddle:develop Apr 1, 2021

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor and simplify hook design & add Tensor.register_hook API #31775

Refactor and simplify hook design & add Tensor.register_hook API #31775

chenwhql commented Mar 22, 2021 •

edited

Loading

paddle-bot-old bot commented Mar 22, 2021

JiabinYang left a comment

JiabinYang Mar 29, 2021

chenwhql Mar 29, 2021

JiabinYang Mar 29, 2021

chenwhql Mar 29, 2021

JiabinYang Mar 29, 2021

chenwhql Mar 29, 2021

JiabinYang Mar 29, 2021

chenwhql Mar 29, 2021

JiabinYang Mar 29, 2021

chenwhql Mar 29, 2021

JiabinYang Mar 29, 2021

chenwhql Mar 29, 2021

JiabinYang Mar 29, 2021

chenwhql Mar 29, 2021

ForFishes left a comment

zhwesky2010 Mar 30, 2021

chenwhql Mar 30, 2021

zhwesky2010 Mar 30, 2021

zhwesky2010 Mar 30, 2021

chenwhql Mar 30, 2021

zhwesky2010 Mar 30, 2021

TCChenlong left a comment

JiabinYang left a comment

Refactor and simplify hook design & add Tensor.register_hook API #31775

Refactor and simplify hook design & add Tensor.register_hook API #31775

Conversation

chenwhql commented Mar 22, 2021 • edited Loading

PR types

PR changes

Describe

1. Refactor

2. Add Tensor.register_hook method

3. Doc

paddle-bot-old bot commented Mar 22, 2021

JiabinYang left a 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

Choose a reason for hiding this comment

ForFishes left a 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

JiabinYang left a comment

Choose a reason for hiding this comment

chenwhql commented Mar 22, 2021 •

edited

Loading