82 lines
2.9 KiB
C++
82 lines
2.9 KiB
C++
// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#pragma once
|
|
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <utility>
|
|
#include <vector>
|
|
namespace paddle {
|
|
namespace imperative {
|
|
|
|
class VariableWrapper;
|
|
|
|
/** [ VariableWrapper Hook ]
|
|
*
|
|
* @brief This hook functor is executed before the grad OpBase is executed or
|
|
* after gradient accumulation completed in current batch.
|
|
* 1. For interior var, VariableWrapper Hook take the input of the
|
|
* current grad OpBase as input.
|
|
* 2. For leaf var, VariableWrapper Hook take the inner_var_ of
|
|
* GradientAccumulator as input.
|
|
*
|
|
* @note This hook functor will not change the input gradient VariableWrapper,
|
|
* but if you copy the input VariableWrapper and change the value of
|
|
* Variable in VariableWrapper, the value of input will also be changed,
|
|
* because they shared same PlaceHolder.
|
|
*
|
|
* @note [ Why need to be OpBase `PreHook`, why not `PostHook`? ]
|
|
*
|
|
* We expect If set OpBase post hook, when the op executed end, the
|
|
* op's output gradient may not be the final state, because it may need
|
|
* other op's gradient output to accumulated to it. But before op can
|
|
* be executed, the gradient output must have been accumulated to final
|
|
* value.
|
|
*
|
|
* @note [ Why Leaf gradient is special? ]
|
|
*
|
|
* Because the leaf VarBase's GradVarBase has no GradOpNode, so leaf
|
|
* GradVarBase has no next OpBase to executed, so if need to deal with
|
|
* the leaf GradVarBase, we should call hooks after gradient accumulation
|
|
* completed.
|
|
*/
|
|
class VariableWrapperHook {
|
|
public:
|
|
virtual ~VariableWrapperHook() = default;
|
|
virtual std::shared_ptr<VariableWrapper> operator()(
|
|
const std::shared_ptr<VariableWrapper>& var) = 0;
|
|
};
|
|
|
|
class CppVariableWrapperHook : public VariableWrapperHook {
|
|
public:
|
|
explicit CppVariableWrapperHook(
|
|
std::function<std::shared_ptr<VariableWrapper>(
|
|
const std::shared_ptr<VariableWrapper>&)>&& fn)
|
|
: fn_(std::move(fn)) {}
|
|
|
|
std::shared_ptr<VariableWrapper> operator()(
|
|
const std::shared_ptr<VariableWrapper>& var) override {
|
|
return fn_(var);
|
|
}
|
|
|
|
private:
|
|
std::function<std::shared_ptr<VariableWrapper>(
|
|
const std::shared_ptr<VariableWrapper>&)>
|
|
fn_;
|
|
};
|
|
|
|
} // namespace imperative
|
|
} // namespace paddle
|