Files
paddlepaddle--paddle/paddle/fluid/imperative/hooks.h
T
2026-07-13 12:40:42 +08:00

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