huawei
/
mindspore2022
0
0
Fork 53
Code Issues Pull Requests 29 Projects Releases Wiki Activity

!30237 Implement forward hook and enhance backward hook for PyNative 

Merge pull request !30237 from JoyLvliang/Implement_forward_hook_for_PyNative
This commit is contained in:
i-robot 2022-02-23 08:38:43 +00:00 committed by Gitee
parent 947ef2af94 b3ff6043ae
commit 35cdaa8118
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
20 changed files with 965 additions and 271 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

74
mindspore/ccsrc/backend/graph_compiler/backend.cc
View File

@ -580,6 +580,30 @@ void MindRTBackend::CompileGraph(const GraphSegmentPtr &segment) {
}
namespace {
ValuePtr GetControlOpInputFromMakeTuple(const std::shared_ptr<GraphCompiler> &graph_compiler,
const AnfNodePtr &front_cnode, const CNodePtr &backend_cnode,
const std::map<KernelWithIndex, tensor::TensorPtr> &op_output_map,
const std::map<AnfNodePtr, size_t> &parameter_index,
const std::vector<tensor::TensorPtr> &graph_inputs,
InputTensorInfo *input_tensor_info, size_t *input_index) {
MS_EXCEPTION_IF_NULL(graph_compiler);
MS_EXCEPTION_IF_NULL(front_cnode);
MS_EXCEPTION_IF_NULL(input_index);
MS_LOG(DEBUG) << "The input node of hook op: " << front_cnode->DebugString() << " is a make tuple node.";
auto make_tuple = front_cnode->cast<CNodePtr>();
MS_EXCEPTION_IF_NULL(make_tuple);
const auto output_size = make_tuple->size() - 1;
std::vector<ValuePtr> output_values;
for (size_t idx = 0; idx < output_size; ++idx) {
TensorPtr tensor = graph_compiler->GetSingleOpInputTensorByIndex(backend_cnode, op_output_map, parameter_index,
graph_inputs, input_tensor_info, *input_index);
MS_EXCEPTION_IF_NULL(tensor);
output_values.emplace_back(tensor);
++(*input_index);
}
return std::make_shared<ValueTuple>(output_values);
}
void GetControlOpInput(const std::shared_ptr<GraphCompiler> &graph_compiler, const CNodePtr &front_cnode,
const CNodePtr &backend_cnode, const std::map<KernelWithIndex, tensor::TensorPtr> &op_output_map,
const std::map<AnfNodePtr, size_t> &parameter_index,
@ -594,34 +618,36 @@ void GetControlOpInput(const std::shared_ptr<GraphCompiler> &graph_compiler, con
for (size_t i = 1; i < inputs.size(); i++) {
const auto &input_node = inputs[i];
MS_EXCEPTION_IF_NULL(input_node);
auto kernel_with_index = AnfAlgo::VisitKernel(input_node, 0);
auto real_input = kernel_with_index.first;
MS_EXCEPTION_IF_NULL(real_input);
if (!real_input->isa<ValueNode>()) {
TensorPtr tensor = graph_compiler->GetSingleOpInputTensorByIndex(backend_cnode, op_output_map, parameter_index,
graph_inputs, input_tensor_info, input_index);
MS_EXCEPTION_IF_NULL(tensor);
args->emplace_back(tensor);
input_index++;
if (IsPrimitiveCNode(input_node, prim::kPrimMakeTuple)) {
// Hook multi-input or multi-output.
args->emplace_back(GetControlOpInputFromMakeTuple(graph_compiler, input_node, backend_cnode, op_output_map,
parameter_index, graph_inputs, input_tensor_info,
&input_index));
continue;
}
// Get value from value node.
const auto &value_node = real_input->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(value_node);
const auto &value = value_node->value();
MS_EXCEPTION_IF_NULL(value);
if (value->isa<ValueSequence>()) {
const auto &value_sequeue = value->cast<ValueSequencePtr>();
MS_EXCEPTION_IF_NULL(value_sequeue);
input_index += value_sequeue->size();
// Hook single-input or single-output.
auto real_input = AnfAlgo::VisitKernel(input_node, 0).first;
MS_EXCEPTION_IF_NULL(real_input);
if (!real_input->isa<ValueNode>()) {
auto tensor = graph_compiler->GetSingleOpInputTensorByIndex(backend_cnode, op_output_map, parameter_index,
graph_inputs, input_tensor_info, input_index);
MS_EXCEPTION_IF_NULL(tensor);
args->emplace_back(tensor);
++input_index;
} else {
input_index++;
const auto &value_node = real_input->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(value_node);
const auto &value = value_node->value();
MS_EXCEPTION_IF_NULL(value);
args->emplace_back(value);
if (value->isa<ValueSequence>()) {
const auto &value_sequeue = value->cast<ValueSequencePtr>();
MS_EXCEPTION_IF_NULL(value_sequeue);
input_index += value_sequeue->size();
} else {
++input_index;
}
}
args->emplace_back(value);
}
}

75
mindspore/ccsrc/frontend/optimizer/ad/dfunctor.cc
View File

@ -186,6 +186,52 @@ static bool HasSideEffectBackProp(const CNodePtr &cnode) {
return false;
}
static AnfNodePtr SkipHookNodeInBackProp(const AnfNodePtr &node) {
MS_EXCEPTION_IF_NULL(node);
if (IsPrimitiveCNode(node, prim::kPrimHookBackward) || IsPrimitiveCNode(node, prim::kPrimCellBackwardHook)) {
MS_LOG(WARNING)
<< "Hook operation does not work in graph mode or ms_function, it will be eliminated during compilation.";
auto output_cnode = node->cast<CNodePtr>();
if (output_cnode->size() - 1 == 1) {
return output_cnode->input(1);
}
// Replace hook node with make tuple node.
abstract::AbstractBasePtrList multi_output_abs;
std::vector<AnfNodePtr> multi_output_nodes{NewValueNode(prim::kPrimMakeTuple)};
std::for_each(output_cnode->inputs().begin() + 1, output_cnode->inputs().end(),
[&multi_output_nodes, &multi_output_abs](const AnfNodePtr &inp) {
MS_EXCEPTION_IF_NULL(inp);
multi_output_nodes.emplace_back(inp);
multi_output_abs.emplace_back(inp->abstract());
});
auto primal_graph = node->func_graph();
MS_EXCEPTION_IF_NULL(primal_graph);
auto make_tuple = primal_graph->NewCNode(std::move(multi_output_nodes));
make_tuple->set_abstract(std::make_shared<abstract::AbstractTuple>(multi_output_abs));
auto mng = primal_graph->manager();
MS_EXCEPTION_IF_NULL(mng);
if (!mng->Replace(node, make_tuple)) {
MS_LOG(EXCEPTION) << "Failed to replace old node: " << node->DebugString()
<< " with new node: " << make_tuple->DebugString();
}
return make_tuple;
}
if (IsPrimitiveCNode(node, prim::kPrimTupleGetItem)) {
auto tuple_get_item = node->cast<CNodePtr>();
auto inp = tuple_get_item->input(1);
if (IsPrimitiveCNode(inp, prim::kPrimHookBackward) || IsPrimitiveCNode(inp, prim::kPrimCellBackwardHook)) {
MS_LOG(WARNING)
<< "Hook operation does not work in graph mode or ms_function, it will be eliminated during compilation.";
constexpr size_t idx = 2;
auto v_node = tuple_get_item->input(idx)->cast<ValueNodePtr>();
MS_EXCEPTION_IF_NULL(v_node);
auto out_idx = GetValue<int64_t>(v_node->value());
return inp->cast<CNodePtr>()->input(out_idx + 1);
}
}
return node;
}
AnfNodePtr HandleRealToComplex(const AnfNodePtr &input, const CNodePtr &din, const FuncGraphPtr &fg) {
MS_EXCEPTION_IF_NULL(input);
TypePtr input_type = input->Type();
@ -244,12 +290,7 @@ void DFunctor::BackPropagate(const CNodePtr &cnode_morph, const CNodePtr &k_app,
}
for (size_t i = 0; i < cnode_morph->size(); i++) {
auto din = tape_->NewCNode({NewValueNode(prim::kPrimTupleGetItem), bprop_app, NewValueNode(SizeToLong(i))});
auto input = cnode_morph->input(i);
// Skip HookBackward op
if (IsPrimitiveCNode(input, prim::kPrimHookBackward)) {
auto inp_i = input->cast<CNodePtr>();
input = inp_i->input(1);
}
auto input = SkipHookNodeInBackProp(cnode_morph->input(i));
auto din_with_real = HandleRealToComplex(input, din, tape_);
MS_EXCEPTION_IF_NULL(din_with_real);
din = din_with_real->cast<CNodePtr>();
@ -299,14 +340,7 @@ AdjointPtr DFunctor::MapMorphism(const AnfNodePtr &morph) {
std::vector<AnfNodePtr> inputs;
std::vector<AdjointPtr> param_adjoints;
for (size_t i = 0; i < cnode_morph->size(); i++) {
auto node = cnode_morph->input(i);
// Skip HookBackward op
if (IsPrimitiveCNode(node, prim::kPrimHookBackward)) {
auto input_i = node->cast<CNodePtr>();
MS_LOG(WARNING)
<< "Hook operation does not work in graph mode or ms_function, it will be eliminated during compilation.";
node = input_i->input(1);
}
auto node = SkipHookNodeInBackProp(cnode_morph->input(i));
AdjointPtr node_adjoint = nullptr;
auto node_adjoint_iter = anfnode_to_adjoin_.find(node);
if (node_adjoint_iter != anfnode_to_adjoin_.end()) {
@ -457,14 +491,9 @@ void DFunctor::MapMorphism() {
// Handle free morphism before output, because in some case, free morphism might depend on output's fv tangent
MapFreeMorphism();
// Skip HookBackward when it is the output node.
// Skip HookBackward op and CellBackwardHook op when it is the output node.
auto output_node = primal_graph_->output();
if (IsPrimitiveCNode(output_node, prim::kPrimHookBackward)) {
auto output_cnode = output_node->cast<CNodePtr>();
MS_LOG(WARNING)
<< "Hook operation does not work in graph mode or ms_function, it will be eliminated during compilation.";
output_node = output_cnode->input(1);
}
output_node = SkipHookNodeInBackProp(output_node);
// Handle morphism from output.
(void)MapMorphism(output_node);
@ -662,7 +691,9 @@ void DFunctor::MapValueObject() {
if (IsValueNode<Primitive>(node)) { // Primitive.
auto prim = GetValueNode<PrimitivePtr>(node);
if (GetValueNode<PrimitivePtr>(node) == prim::kPrimReturn ||
(prim->Hash() == prim::kPrimHookBackward->Hash() && prim->name() == prim::kPrimHookBackward->name())) {
(prim->Hash() == prim::kPrimHookBackward->Hash() && prim->name() == prim::kPrimHookBackward->name()) ||
(prim->Hash() == prim::kPrimCellBackwardHook->Hash() &&
prim->name() == prim::kPrimCellBackwardHook->name())) {
continue;
}
MS_LOG(DEBUG) << "Map Primitive node " << node->DebugString() << ".";

2
mindspore/ccsrc/frontend/optimizer/ad/dfunctor.h
View File

@ -159,6 +159,8 @@ class KPrim {
private:
FuncGraphPtr GetBprop(const PrimitivePtr &prim, const pipeline::ResourceBasePtr &resources = nullptr);
FuncGraphPtr GetFprop(const PrimitivePtr &prim);
FuncGraphPtr GetPrimBprop(const PrimitivePtr &prim, const ValueNodePtr &value_node,
const pipeline::ResourceBasePtr &resources);
FuncGraphPtr FakeBprop(const ValueNodePtr &value_node, const pipeline::ResourceBasePtr &resources);
FuncGraphPtr BpropCut(const ValueNodePtr &value_node, const pipeline::ResourceBasePtr &resources);
// Given a bprop rule, do the K mapping.

46
mindspore/ccsrc/frontend/optimizer/ad/kprim.cc
View File

@ -335,6 +335,28 @@ void KPrim::ExportBpropMindir(const py::object &obj) {
}
#endif
FuncGraphPtr KPrim::GetPrimBprop(const PrimitivePtr &prim, const ValueNodePtr &value_node,
const pipeline::ResourceBasePtr &resources) {
MS_EXCEPTION_IF_NULL(prim);
MS_EXCEPTION_IF_NULL(value_node);
FuncGraphPtr bprop_fg = nullptr;
auto iter = bprop_registry_.find(prim);
if (iter != bprop_registry_.end()) {
bprop_fg = iter->second;
}
if (bprop_fg == nullptr) {
bprop_fg = GetBprop(prim, resources);
if (bprop_fg != nullptr) {
// Set bprop_g graph cache
bprop_registry_[prim] = bprop_fg;
} else {
bprop_fg = FakeBprop(value_node, resources);
}
}
return bprop_fg;
}
FuncGraphPtr KPrim::GetBprop(const PrimitivePtr &prim, const pipeline::ResourceBasePtr &resources) {
// Set a child scope named "grad'PrimitiveName'" for the bprop function,
// and add "Gradients" to the front.
@ -518,28 +540,16 @@ FuncGraphPtr KPrim::KPrimitive(const CNodePtr &cnode, const ValueNodePtr &value_
}
FuncGraphPtr bprop_fg = nullptr;
if (prim->Hash() == prim::kPrimHookBackward->Hash() && prim->name() == prim::kPrimHookBackward->name()) {
if ((prim->Hash() == prim::kPrimHookBackward->Hash() && prim->name() == prim::kPrimHookBackward->name()) ||
(prim->Hash() == prim::kPrimCellBackwardHook->Hash() && prim->name() == prim::kPrimCellBackwardHook->name())) {
if (MsContext::GetInstance()->get_param<int>(MsCtxParam::MS_CTX_EXECUTION_MODE) == kGraphMode) {
MS_LOG(EXCEPTION)
<< "The Primitive 'HookBackward' is not supported in graph mode, which is only supported in pynative mode.\n"
<< "The Hook operation is not supported in graph mode, which is only supported in pynative mode.\n"
<< trace::GetDebugInfo(cnode->debug_info());
}
bprop_fg = BpropCut(value_node, resources);
} else {
auto iter = bprop_registry_.find(prim);
if (iter != bprop_registry_.end()) {
bprop_fg = iter->second;
}
if (bprop_fg == nullptr) {
bprop_fg = GetBprop(prim, resources);
if (bprop_fg != nullptr) {
// Set bprop_g graph cache
bprop_registry_[prim] = bprop_fg;
} else {
bprop_fg = FakeBprop(value_node, resources);
}
}
bprop_fg = GetPrimBprop(prim, value_node, resources);
}
AdjustForAutoMonad(prim, bprop_fg);
@ -766,8 +776,10 @@ FuncGraphPtr KPrim::BpropCut(const ValueNodePtr &value_node, const pipeline::Res
auto func_graph = std::make_shared<FuncGraph>();
std::vector<AnfNodePtr> outputs;
auto prim_py = prim->cast<PrimitivePyPtr>();
MS_EXCEPTION_IF_NULL(prim_py);
auto bprop_cut = std::make_shared<PrimitivePy>("bprop_cut");
bprop_cut->CopyHookFunction(prim);
bprop_cut->CopyHookFunction(prim_py);
auto cell_id = GetValue<std::string>(prim->GetAttr("cell_id"));
if (cell_id != "") {

7
mindspore/ccsrc/frontend/optimizer/ad/kpynative.cc
View File

@ -398,7 +398,7 @@ bool KPynativeCellImpl::KPynativeOp(const CNodePtr &cnode, const ValuePtrList &o
}
FuncGraphPtr bprop_fg = nullptr;
if (IsPrimitiveEquals(prim, prim::kPrimHookBackward)) {
if (IsPrimitiveEquals(prim, prim::kPrimHookBackward) || IsPrimitiveEquals(prim, prim::kPrimCellBackwardHook)) {
bprop_fg = BuildBPropCutFuncGraph(prim, cnode);
} else if (IsPrimitiveEquals(prim, prim::kPrimMakeTuple) || IsPrimitiveEquals(prim, prim::kPrimMakeList)) {
bprop_fg = BuildMakeSequenceBprop(prim, cnode);
@ -906,8 +906,11 @@ FuncGraphPtr KPynativeCellImpl::BuildBPropCutFuncGraph(const PrimitivePtr &prim,
auto func_graph = std::make_shared<FuncGraph>();
std::vector<AnfNodePtr> outputs;
auto prim_py = prim->cast<PrimitivePyPtr>();
MS_EXCEPTION_IF_NULL(prim_py);
auto bprop_cut = std::make_shared<PrimitivePy>("bprop_cut");
bprop_cut->CopyHookFunction(prim);
bprop_cut->CopyHookFunction(prim_py);
prim_py->AddBpropCutPrim(bprop_cut);
auto cell_id = GetValue<std::string>(prim->GetAttr("cell_id"));
if (cell_id != "") {

3
mindspore/ccsrc/frontend/optimizer/ad/prim_bprop_optimizer.cc
View File

@ -194,7 +194,8 @@ FuncGraphPtr PrimBpropOptimizer::OptimizeBPropFuncGraph(const FuncGraphPtr &bpro
MS_LOG(DEBUG) << "Hash of prim " << prim->ToString() << " is:" << prim->hash();
// kPrimHookBackward
bool hookback_flg = IsPrimitiveEquals(prim, prim::kPrimHookBackward);
bool hookback_flg =
IsPrimitiveEquals(prim, prim::kPrimHookBackward) || IsPrimitiveEquals(prim, prim::kPrimCellBackwardHook);
if (hookback_flg || IsPrimitiveEquals(prim, prim::kPrimMakeTuple) || IsPrimitiveEquals(prim, prim::kPrimMakeList)) {
return GenSpecOptBprop(bprop_fg, op_args, out, prim, hookback_flg);
}

8
mindspore/ccsrc/frontend/optimizer/irpass.cc
View File

@ -61,10 +61,10 @@ OptimizeIRPassLib::OptimizeIRPassLib() {
prim::kPrimIdentity, prim::kPrimMomentum, prim::kPrimMul, prim::kPrimPow});
arithmetic_simplify2_ =
MakeSubstitution(std::make_shared<ArithmeticSimplify2>(), "arithmetic_simplify2", {prim::kPrimMul});
special_op_eliminate_ =
MakeSubstitution(std::make_shared<SpecialOpEliminater>(), "special_op_eliminate",
{prim::kPrimInsertGradientOf, prim::kPrimStopGradient, prim::kPrimHookBackward,
prim::kPrimPrintShapeType, prim::kPrimGetRefValue, prim::kPrimMirror, prim::kPrimVirtualDiv});
special_op_eliminate_ = MakeSubstitution(
std::make_shared<SpecialOpEliminater>(), "special_op_eliminate",
{prim::kPrimInsertGradientOf, prim::kPrimStopGradient, prim::kPrimHookBackward, prim::kPrimCellBackwardHook,
prim::kPrimPrintShapeType, prim::kPrimGetRefValue, prim::kPrimMirror, prim::kPrimVirtualDiv});
pynative_eliminate_ = MakeSubstitution(std::make_shared<PynativeEliminater>(), "pynative_eliminate", IsCNodeDup);
zero_like_fill_zero_ =
MakeSubstitution(std::make_shared<ZeroLikeFillZero>(), "zero_like_fill_zero", prim::kPrimZerosLike);

8
mindspore/ccsrc/frontend/optimizer/irpass/special_op_eliminate.h
View File

@ -45,6 +45,7 @@ class SpecialOpEliminater : public OptimizerCaller {
: insert_gradient_of_(std::make_shared<PrimEliminater>(prim::kPrimInsertGradientOf)),
stop_gradient_(std::make_shared<PrimEliminater>(prim::kPrimStopGradient)),
hook_backward_(std::make_shared<PrimEliminater>(prim::kPrimHookBackward)),
cell_backward_hook_(std::make_shared<PrimEliminater>(prim::kPrimCellBackwardHook)),
print_shape_type_(std::make_shared<PrimEliminater>(prim::kPrimPrintShapeType)),
get_ref_value_(std::make_shared<PrimEliminater>(prim::kPrimGetRefValue)),
mirror_(std::make_shared<PrimEliminater>(prim::kPrimMirror)),
@ -52,6 +53,7 @@ class SpecialOpEliminater : public OptimizerCaller {
eliminaters_.emplace_back(insert_gradient_of_);
eliminaters_.emplace_back(stop_gradient_);
eliminaters_.emplace_back(hook_backward_);
eliminaters_.emplace_back(cell_backward_hook_);
eliminaters_.emplace_back(print_shape_type_);
eliminaters_.emplace_back(get_ref_value_);
eliminaters_.emplace_back(mirror_);
@ -64,7 +66,7 @@ class SpecialOpEliminater : public OptimizerCaller {
for (auto &eliminater : eliminaters_) {
new_node = (*eliminater)(optimizer, node);
if (new_node != nullptr) {
if (IsPrimitiveCNode(node, prim::kPrimHookBackward)) {
if (IsPrimitiveCNode(node, prim::kPrimHookBackward) || IsPrimitiveCNode(node, prim::kPrimCellBackwardHook)) {
MS_LOG(WARNING)
<< "Hook operation does not work in graph mode or ms_function, it will be eliminated during compilation.";
}
@ -75,8 +77,8 @@ class SpecialOpEliminater : public OptimizerCaller {
}
private:
OptimizerCallerPtr insert_gradient_of_, stop_gradient_, hook_backward_, print_shape_type_, get_ref_value_, mirror_,
virtual_div_;
OptimizerCallerPtr insert_gradient_of_, stop_gradient_, hook_backward_, cell_backward_hook_, print_shape_type_,
get_ref_value_, mirror_, virtual_div_;
std::vector<OptimizerCallerPtr> eliminaters_{};
};

2
mindspore/ccsrc/pipeline/jit/parse/data_converter.cc
View File

@ -141,7 +141,7 @@ FuncGraphPtr ConvertToBpropCut(const py::object &obj) {
std::vector<AnfNodePtr> outputs;
auto fake_bprop = std::make_shared<PrimitivePy>("bprop_cut");
fake_bprop->set_hook(bprop_func);
fake_bprop->AddBackwardHookFn(0, bprop_func);
(void)fake_bprop->AddAttr(CUSTOM_BPROP_NAME, MakeValue(true));
outputs.push_back(NewValueNode(fake_bprop));

46
mindspore/ccsrc/pipeline/pynative/pynative_execute.cc
View File

@ -75,8 +75,8 @@ const size_t ARG_SIZE = 2;
const size_t MAX_TOP_CELL_COUNTS = 20;
// primitive unable to infer value for constant input in PyNative mode
const std::set<std::string> kVmOperators = {"make_ref", "HookBackward", "InsertGradientOf", "stop_gradient",
"mixed_precision_cast"};
const std::set<std::string> kVmOperators = {"make_ref", "InsertGradientOf", "stop_gradient", "mixed_precision_cast",
"HookBackward", "CellBackwardHook"};
const char kOpsFunctionModelName[] = "mindspore.ops.functional";
const char kGrad[] = "grad";
std::map<std::string, std::shared_ptr<session::SessionBasic>> kSessionBackends;
@ -1728,6 +1728,22 @@ void GradExecutor::EnableOpGraphCache(bool is_enable) {
inst->set_param<bool>(MS_CTX_ENABLE_PYNATIVE_OP_GRAPH_CACHE, is_enable);
}
void GradExecutor::SetHookChanged(const py::object &cell) {
auto cell_id = GetId(cell);
for (const auto &top_cell : top_cell_list_) {
MS_EXCEPTION_IF_NULL(top_cell);
if (top_cell->cell_id().find(cell_id) != std::string::npos) {
top_cell->set_hook_changed(true);
}
const auto &sub_cells = top_cell->sub_cell_list();
for (const auto &sub_cell_id : sub_cells) {
if (sub_cell_id.find(cell_id) != std::string::npos) {
top_cell->set_hook_changed(true);
}
}
}
}
void GradExecutor::RecordGradOpInfo(const OpExecInfoPtr &op_exec_info) {
if (!grad_flag_) {
MS_LOG(DEBUG) << "Grad flag is set to false, no need to record op info";
@ -2050,14 +2066,17 @@ py::object ForwardExecutor::RunOpInVM(const OpExecInfoPtr &op_exec_info) {
MS_EXCEPTION_IF_NULL(op_exec_info->py_primitive);
auto &op_inputs = op_exec_info->op_inputs;
if (op_exec_info->op_name == "HookBackward" || op_exec_info->op_name == "InsertGradientOf" ||
op_exec_info->op_name == "stop_gradient") {
if (op_exec_info->op_name == prim::kPrimInsertGradientOf->name() ||
op_exec_info->op_name == prim::kPrimStopGradient->name() ||
op_exec_info->op_name == prim::kPrimHookBackward->name() ||
op_exec_info->op_name == prim::kPrimCellBackwardHook->name()) {
py::tuple result(op_inputs.size());
for (size_t i = 0; i < op_inputs.size(); i++) {
py::object input = op_inputs[i];
auto tensor = py::cast<tensor::TensorPtr>(input);
MS_EXCEPTION_IF_NULL(tensor);
if (op_exec_info->op_name == "HookBackward") {
if (op_exec_info->op_name == prim::kPrimHookBackward->name() ||
op_exec_info->op_name == prim::kPrimCellBackwardHook->name()) {
// the input object is not a output of forward cnode, eg: parameter
result[i] = tensor;
} else {
@ -2405,7 +2424,10 @@ void GradExecutor::NewGraphInner(py::object *ret, const py::object &cell, const
// Top cell forward run.
const auto &pre_top_cell = top_it->second;
MS_EXCEPTION_IF_NULL(pre_top_cell);
if (!pre_top_cell->is_dynamic()) {
if (pre_top_cell->hook_changed()) {
already_run_top_cell_.erase(top_it);
EraseTopCellFromTopCellList(pre_top_cell);
} else if (!pre_top_cell->is_dynamic()) {
MS_LOG(DEBUG) << "Top cell " << cell_id << " is not dynamic, no need to run NewGraphInner again";
ResetTopCellInfo(pre_top_cell, args);
PushHighOrderGraphStack(pre_top_cell);
@ -2594,7 +2616,7 @@ void GradExecutor::DoGradForCustomBprop(const py::object &cell, const py::object
auto cell_ptr = py::cast<CellPtr>(cell);
fake_prim->set_bprop_cls_name(cell_ptr->name());
}
fake_prim->set_hook(bprop_func);
fake_prim->AddBackwardHookFn(0, bprop_func);
const auto &cell_id = GetCellId(cell, args);
(void)fake_prim->AddAttr("cell_id", MakeValue(cell_id));
@ -2911,7 +2933,7 @@ py::object GradExecutor::CheckAlreadyRun(const prim::GradOperationPtr &grad, con
auto find_top_cell = GetTopCell(check_already_run_cell_id);
if (find_top_cell != nullptr) {
MS_LOG(DEBUG) << "Find already run top cell";
forward_run = find_top_cell->forward_already_run();
forward_run = find_top_cell->forward_already_run() && !find_top_cell->hook_changed();
auto curr_top_cell = top_cell();
set_top_cell(find_top_cell);
bool input_args_changed =
@ -3293,6 +3315,13 @@ bool PynativeExecutor::grad_flag() const { return grad_executor()->grad_flag();
void PynativeExecutor::set_grad_flag(bool flag) { grad_executor()->set_grad_flag(flag); }
void PynativeExecutor::SetHookChanged(const py::object &cell) {
if (!py::isinstance<Cell>(cell)) {
MS_LOG(EXCEPTION) << "The 'set_hook_changed' function is only supported on Cell object!";
}
grad_executor()->SetHookChanged(cell);
}
void PynativeExecutor::set_graph_phase(const std::string &graph_phase) {
grad_executor()->set_graph_phase(graph_phase);
}
@ -3475,6 +3504,7 @@ REGISTER_PYBIND_DEFINE(PynativeExecutor_, ([](const py::module *m) {
.def("__call__", &PynativeExecutor::Run, "pynative executor run grad graph.")
.def("set_graph_phase", &PynativeExecutor::set_graph_phase, "pynative set graph phase")
.def("grad_flag", &PynativeExecutor::grad_flag, "pynative grad flag")
.def("set_hook_changed", &PynativeExecutor::SetHookChanged, "set pynative hook changed")
.def("set_grad_position", &PynativeExecutor::set_grad_position, "set pynative grad position")
.def("set_grad_flag", &PynativeExecutor::set_grad_flag, py::arg("flag") = py::bool_(false),
"Executor set grad flag.")

5
mindspore/ccsrc/pipeline/pynative/pynative_execute.h
View File

@ -78,6 +78,8 @@ class TopCellInfo {
void set_grad_order(size_t grad_order) { grad_order_ = grad_order; }
bool is_dynamic() const { return is_dynamic_; }
void set_is_dynamic(bool is_dynamic) { is_dynamic_ = is_dynamic; }
bool hook_changed() const { return hook_changed_; }
void set_hook_changed(bool hook_changed) { hook_changed_ = hook_changed; }
bool vm_compiled() const { return vm_compiled_; }
void set_vm_compiled(bool vm_compiled) { vm_compiled_ = vm_compiled; }
bool ms_function_flag() const { return ms_function_flag_; }
@ -123,6 +125,7 @@ class TopCellInfo {
bool is_topest_{false};
bool is_dynamic_{false};
bool vm_compiled_{false};
bool hook_changed_{false};
bool ms_function_flag_{false};
bool is_init_kpynative_{false};
bool forward_already_run_{false};
@ -191,6 +194,7 @@ class GradExecutor {
size_t GetHighOrderStackSize() const { return high_order_stack_.size(); }
TopCellInfoPtr GetTopCell(const string &already_run_cell_id);
void EnableOpGraphCache(bool is_enable);
void SetHookChanged(const py::object &cell);
bool need_renormalize() const { return need_renormalize_; }
bool enable_op_cache() const { return enable_op_cache_; }
bool grad_is_running() const { return grad_is_running_; }
@ -397,6 +401,7 @@ class PynativeExecutor : public std::enable_shared_from_this<PynativeExecutor> {
void set_graph_phase(const std::string &graph_phase);
void set_py_exe_path(const py::object &py_exe_path);
void set_kernel_build_server_dir(const py::object &kernel_build_server_dir);
void SetHookChanged(const py::object &cell);
void NewGraph(const py::object &cell, const py::args &args);
void EndGraph(const py::object &cell, const py::object &out, const py::args &args);
void GradNet(const prim::GradOperationPtr &grad, const py::object &cell, const py::object &weights,

298
mindspore/ccsrc/pybind_api/ir/primitive_py.cc
View File

@ -55,6 +55,54 @@ void SyncData(const py::object &arg) {
tensor->data_sync();
}
}
void ConvertCTensorToPyTensor(const py::tuple &input_args, py::tuple *convert_args) {
MS_EXCEPTION_IF_NULL(convert_args);
if (input_args.size() != (*convert_args).size()) {
MS_LOG(EXCEPTION) << "The size of input_args: " << input_args.size()
<< " should be equal to the size of convert_args: " << (*convert_args).size();
}
for (size_t i = 0; i < input_args.size(); ++i) {
if (py::isinstance<tensor::Tensor>(input_args[i])) {
(*convert_args)[i] = parse::python_adapter::CallPyFn(parse::PYTHON_MOD_PARSE_MODULE,
parse::PYTHON_MOD_CONVERT_TO_MS_TENSOR, input_args[i]);
} else if (py::isinstance<py::tuple>(input_args[i])) {
auto tuple_inp_arg = py::cast<py::tuple>(input_args[i]);
py::tuple convert_tuple_arg(tuple_inp_arg.size());
ConvertCTensorToPyTensor(tuple_inp_arg, &convert_tuple_arg);
(*convert_args)[i] = convert_tuple_arg;
} else {
(*convert_args)[i] = input_args[i];
}
}
}
py::tuple ConstructCellHookFnArgs(const std::string &cell_id, const py::object &grad_input,
const py::object &grad_output) {
constexpr size_t grad_input_index = 1;
constexpr size_t grad_output_index = 2;
constexpr size_t input_args_nums = 3;
// Convert c++ object to python object.
py::tuple c_grad_args(input_args_nums - 1);
c_grad_args[0] = grad_input;
c_grad_args[1] = grad_output;
py::tuple py_grad_args(input_args_nums - 1);
ConvertCTensorToPyTensor(c_grad_args, &py_grad_args);
// Get tuple args of cell hook function.
py::tuple hook_fn_args(input_args_nums);
hook_fn_args[0] = cell_id;
if (!py::isinstance<py::tuple>(py_grad_args[0])) {
hook_fn_args[grad_input_index] = py::make_tuple(py_grad_args[0]);
} else {
hook_fn_args[grad_input_index] = py_grad_args[0];
}
if (!py::isinstance<py::tuple>(py_grad_args[1])) {
hook_fn_args[grad_output_index] = py::make_tuple(py_grad_args[1]);
} else {
hook_fn_args[grad_output_index] = py_grad_args[1];
}
return hook_fn_args;
}
} // namespace
std::map<std::string, py::object> PrimitivePy::hook_grad_;
@ -68,7 +116,9 @@ PrimitivePy::PrimitivePy(const py::object &python_obj, const PrimitivePyAdapterP
Primitive::set_prim_type(adapter->prim_type_);
Primitive::set_const_prim(adapter->is_const_prim_);
Primitive::set_const_input_indexes(adapter->const_input_indexes_);
set_hook(adapter->hook_);
for (const auto &elem : adapter->backward_hook_fn_) {
AddBackwardHookFn(elem.first, elem.second);
}
set_instance_name(adapter->instance_name_);
}
PrimitivePy::~PrimitivePy() {}
@ -91,7 +141,9 @@ py::function PrimitivePy::GetBpropFunction() {
py::tuple check_bprop_out(const py::object &grads_obj, const py::tuple &py_args, const std::string &bprop_cls_name) {
py::tuple grads;
if (!py::isinstance<py::tuple>(grads_obj)) {
if (py::isinstance<py::none>(grads_obj)) {
MS_EXCEPTION(TypeError) << "The 'grads_obj' is none.";
} else if (!py::isinstance<py::tuple>(grads_obj)) {
grads = py::make_tuple(grads_obj);
} else {
grads = py::cast<py::tuple>(grads_obj);
@ -138,21 +190,37 @@ py::tuple check_bprop_out(const py::object &grads_obj, const py::tuple &py_args,
return grads;
}
void PrimitivePy::ConvertCTensorToPyTensor(const py::tuple &input_args, py::tuple *convert_args) const {
MS_EXCEPTION_IF_NULL(convert_args);
if (input_args.size() != (*convert_args).size()) {
MS_LOG(EXCEPTION) << "The size of input_args: " << input_args.size()
<< " should be equal to the size of convert_args: " << (*convert_args).size();
}
for (size_t i = 0; i < input_args.size(); ++i) {
(*convert_args)[i] = py::isinstance<tensor::Tensor>(input_args[i])
? parse::python_adapter::CallPyFn(parse::PYTHON_MOD_PARSE_MODULE,
parse::PYTHON_MOD_CONVERT_TO_MS_TENSOR, input_args[i])
: input_args[i];
void PrimitivePy::AddBpropCutPrim(const PrimitivePyPtr &bprop_cut_prim) {
MS_EXCEPTION_IF_NULL(bprop_cut_prim);
bprop_cut_prims_.emplace_back(bprop_cut_prim);
}
void PrimitivePy::AddBackwardHookFn(const int &key, const py::function &backward_hook_fn) {
backward_hook_fn_[key] = backward_hook_fn;
for (const auto &elem : bprop_cut_prims_) {
PrimitivePyPtr bprop_cut_prim = elem.lock();
if (bprop_cut_prim != nullptr) {
bprop_cut_prim->AddBackwardHookFn(key, backward_hook_fn);
}
}
}
void PrimitivePy::CheckHookConsistency(const py::object &grad_out, const py::object &expected_grad_out) const {
void PrimitivePy::RemoveBackwardHookFn(const int &key) {
auto iter = backward_hook_fn_.find(key);
if (iter != backward_hook_fn_.end()) {
backward_hook_fn_.erase(key);
}
// Remove hook_fn for bprop cut prim on grad graph.
for (const auto &elem : bprop_cut_prims_) {
PrimitivePyPtr bprop_cut_prim = elem.lock();
if (bprop_cut_prim != nullptr) {
bprop_cut_prim->RemoveBackwardHookFn(key);
}
}
}
void PrimitivePy::CheckHookConsistency(const py::object &grad_out, const py::object &expected_grad_out,
const py::object &code_obj, const py::object &co_name) const {
if (py::isinstance<py::tuple>(expected_grad_out)) {
if (!py::isinstance<py::tuple>(grad_out)) {
hook_grad_.clear();
@ -166,15 +234,13 @@ void PrimitivePy::CheckHookConsistency(const py::object &grad_out, const py::obj
<< ", but it is " << actual_out_tuple.size();
}
for (size_t i = 0; i < expected_out_tuple.size(); ++i) {
CheckHookConsistency(actual_out_tuple[i], expected_out_tuple[i]);
CheckHookConsistency(actual_out_tuple[i], expected_out_tuple[i], code_obj, co_name);
}
}
if (py::isinstance<tensor::Tensor>(expected_grad_out)) {
if (!py::isinstance<tensor::Tensor>(grad_out)) {
hook_grad_.clear();
py::object code_obj = py::getattr(hook_, "__code__");
py::object co_name = py::getattr(code_obj, "co_name");
MS_EXCEPTION(TypeError) << "The output type of:" << py::str(co_name) << " should be a tensor but got "
<< py::cast<std::string>(grad_out.attr("__class__").attr("__name__")) << ".";
}
@ -184,8 +250,6 @@ void PrimitivePy::CheckHookConsistency(const py::object &grad_out, const py::obj
MS_EXCEPTION_IF_NULL(expected_out_tensor);
if (actual_out_tensor->GetShapeAndDataTypeInfo() != expected_out_tensor->GetShapeAndDataTypeInfo()) {
hook_grad_.clear();
py::object code_obj = py::getattr(hook_, "__code__");
py::object co_name = py::getattr(code_obj, "co_name");
MS_EXCEPTION(ValueError) << "The output type of " << py::str(co_name)
<< " is not consistent with the expected, it should be "
<< expected_out_tensor->GetShapeAndDataTypeInfo() << ", but got "
@ -195,24 +259,31 @@ void PrimitivePy::CheckHookConsistency(const py::object &grad_out, const py::obj
}
BaseRef PrimitivePy::RunCellBpropFunction(const py::tuple &py_args) const {
if (backward_hook_fn_.size() > 1) {
MS_LOG(EXCEPTION) << "Multiple registration of bprop function is not supported.";
}
SyncData(py_args);
auto size = py_args.size();
constexpr size_t grad_param_nums = 2;
py::tuple input_args(size - grad_param_nums);
for (size_t i = 0; i < size - grad_param_nums; ++i) {
py::tuple converted_args(py_args.size());
ConvertCTensorToPyTensor(py_args, &converted_args);
constexpr size_t non_inp_args_size = 2; // out and dout.
auto inp_args_size = py_args.size() - non_inp_args_size;
py::tuple input_args(inp_args_size);
for (size_t i = 0; i < inp_args_size; ++i) {
input_args[i] = py_args[i];
}
py::tuple convert_args(py_args.size());
ConvertCTensorToPyTensor(py_args, &convert_args);
// Run bprop function.
auto inst = pynative::PynativeExecutor::GetInstance();
MS_EXCEPTION_IF_NULL(inst);
try {
MS_LOG(DEBUG) << "Run bprop function start";
inst->NewGraph(hook_, input_args.cast<py::args>());
py::object grads_obj = hook_(*convert_args);
py::tuple grads = check_bprop_out(grads_obj, py_args, bprop_cls_name_);
inst->EndGraph(hook_, grads_obj, input_args.cast<py::args>());
MS_LOG(DEBUG) << "Run bprop function end";
MS_LOG(DEBUG) << "Run bprop function start.";
py::tuple grads;
for (const auto &elem : backward_hook_fn_) {
inst->NewGraph(elem.second, input_args.cast<py::args>());
py::object grads_obj = elem.second(*converted_args);
grads = check_bprop_out(grads_obj, py_args, bprop_cls_name_);
inst->EndGraph(elem.second, grads_obj, input_args.cast<py::args>());
}
MS_LOG(DEBUG) << "Run bprop function end.";
return std::make_shared<PyObjectRef>(grads);
} catch (std::exception &bt) {
inst->ClearRes();
@ -221,62 +292,61 @@ BaseRef PrimitivePy::RunCellBpropFunction(const py::tuple &py_args) const {
}
BaseRef PrimitivePy::RunCellHookFunction(const py::tuple &py_args) const {
constexpr size_t grad_input_index = 1;
constexpr size_t grad_output_index = 2;
constexpr size_t input_param_nums = 3;
SyncData(py_args[grad_output_index]);
py::object obj;
// Get the gradient passed to current bprop cut op.
const auto args_size = py_args.size();
py::object grad_output = py_args[args_size - 1];
// Get the cell id.
auto cell_id = GetValue<std::string>(this->GetAttr(kCellIDAttrName));
auto iter = hook_grad_.find(cell_id);
if (iter != hook_grad_.end()) {
py::object code_obj = py::getattr(hook_, "__code__");
py::object co_name = py::getattr(code_obj, "co_name");
if (std::string(py::str(co_name)) == "staging_specialize") {
py::object name_obj = py::getattr(hook_, "__name__");
MS_LOG(EXCEPTION) << "Decorating hook function " << py::str(name_obj) << " with '@ms_function' is not supported.";
// The second bprop_cut used to hook output gradient of cell.
for (const auto &elem : backward_hook_fn_) {
py::object code_obj = py::getattr(elem.second, "__code__");
py::object co_name = py::getattr(code_obj, "co_name");
if (std::string(py::str(co_name)) == "staging_specialize") {
py::object name_obj = py::getattr(elem.second, "__name__");
MS_LOG(EXCEPTION) << "Decorating hook function " << py::str(name_obj)
<< " with '@ms_function' is not supported.";
}
SyncData(grad_output);
py::tuple hook_fn_args = ConstructCellHookFnArgs(cell_id, iter->second, grad_output);
py::object ret = elem.second(*hook_fn_args);
if (!py::isinstance<py::none>(ret)) {
grad_output = ret;
}
CheckHookConsistency(grad_output, py_args[args_size - 1], code_obj, co_name);
}
py::tuple convert_args(input_param_nums - 1);
py::tuple input_args(input_param_nums - 1);
input_args[0] = iter->second;
input_args[1] = py_args[grad_output_index];
ConvertCTensorToPyTensor(input_args, &convert_args);
auto hook_args = py::tuple(input_param_nums);
hook_args[0] = cell_id;
hook_args[grad_input_index] = py::make_tuple(convert_args[0]);
hook_args[grad_output_index] = py::make_tuple(convert_args[1]);
obj = hook_(*hook_args);
if (py::isinstance<py::none>(obj)) {
obj = py_args[grad_output_index];
}
CheckHookConsistency(obj, py_args[grad_output_index]);
(void)hook_grad_.erase(cell_id);
} else {
hook_grad_[cell_id] = py_args[grad_output_index];
obj = py_args[grad_output_index];
// The first bprop_cut used to hook input gradient of cell.
SyncData(grad_output);
hook_grad_[cell_id] = grad_output;
}
obj = py::make_tuple(obj);
return std::make_shared<PyObjectRef>(obj);
if (!py::isinstance<py::tuple>(grad_output)) {
grad_output = py::make_tuple(grad_output);
}
return std::make_shared<PyObjectRef>(grad_output);
}
BaseRef PrimitivePy::RunVariableHookFunction(const py::tuple &py_args) const {
py::object code_obj = py::getattr(hook_, "__code__");
py::object co_name = py::getattr(code_obj, "co_name");
if (std::string(py::str(co_name)) == "staging_specialize") {
py::object name_obj = py::getattr(hook_, "__name__");
MS_LOG(EXCEPTION) << "Decorating hook function " << py::str(name_obj) << " with '@ms_function' is not supported.";
}
constexpr size_t grad_output_index = 2;
SyncData(py_args[grad_output_index]);
py::object obj = hook_(py::make_tuple(py_args[grad_output_index]));
if (py::isinstance<py::none>(obj)) {
obj = py_args[grad_output_index];
py::object grad_output = py_args[grad_output_index];
for (const auto &elem : backward_hook_fn_) {
py::object code_obj = py::getattr(elem.second, "__code__");
py::object co_name = py::getattr(code_obj, "co_name");
if (std::string(py::str(co_name)) == "staging_specialize") {
py::object name_obj = py::getattr(elem.second, "__name__");
MS_LOG(EXCEPTION) << "Decorating hook function " << py::str(name_obj) << " with '@ms_function' is not supported.";
}
SyncData(grad_output);
py::object ret = elem.second(py::make_tuple(grad_output));
if (!py::isinstance<py::none>(ret)) {
grad_output = ret;
}
CheckHookConsistency(grad_output, py_args[grad_output_index], code_obj, co_name);
}
CheckHookConsistency(obj, py_args[grad_output_index]);
obj = py::make_tuple(obj);
return std::make_shared<PyObjectRef>(obj);
grad_output = py::make_tuple(grad_output);
return std::make_shared<PyObjectRef>(grad_output);
}
BaseRef PrimitivePy::RunHookFunction(const VectorRef &args) const {
@ -321,14 +391,12 @@ py::dict PrimitivePy::GetAttrDict() {
return attr_dict;
}
void PrimitivePy::CopyHookFunction(const PrimitivePtr &primitive) {
MS_EXCEPTION_IF_NULL(primitive);
if (!primitive->isa<PrimitivePy>()) {
MS_LOG(EXCEPTION) << "Cannot copy a primitive which is not python primitive hook function to python primitive!";
}
auto primitive_py = primitive->cast<PrimitivePyPtr>();
void PrimitivePy::CopyHookFunction(const PrimitivePyPtr &primitive_py) {
MS_EXCEPTION_IF_NULL(primitive_py);
this->set_hook(primitive_py->hook());
const auto &backward_hook_fn = primitive_py->backward_hook_fn();
for (const auto &elem : backward_hook_fn) {
AddBackwardHookFn(elem.first, elem.second);
}
if (primitive_py->HasAttr(kBpropAttrName)) {
set_bprop_cls_name(primitive_py->bprop_cls_name_);
(void)this->AddAttr(kBpropAttrName, primitive_py->GetAttr(kBpropAttrName));
@ -500,11 +568,24 @@ void PrimitivePyAdapter::set_signatures(const std::vector<Signature> &signatures
}
}
void PrimitivePyAdapter::set_hook(const py::function &hook) {
hook_ = hook;
int PrimitivePyAdapter::AddBackwardHookFn(const py::function &backward_hook_fn) {
++backward_hook_fn_key_;
backward_hook_fn_[backward_hook_fn_key_] = backward_hook_fn;
auto prim = attached_primitive_.lock();
if (prim != nullptr) {
prim->set_hook(hook);
prim->AddBackwardHookFn(backward_hook_fn_key_, backward_hook_fn);
}
return backward_hook_fn_key_;
}
void PrimitivePyAdapter::RemoveBackwardHookFn(int key) {
auto iter = backward_hook_fn_.find(key);
if (iter != backward_hook_fn_.end()) {
backward_hook_fn_.erase(iter);
}
auto prim = attached_primitive_.lock();
if (prim != nullptr) {
prim->RemoveBackwardHookFn(key);
}
}
@ -524,27 +605,28 @@ void PrimitivePyAdapter::set_attached_primitive(const PrimitivePyPtr &prim) {
attached_primitive_ = prim;
}
REGISTER_PYBIND_DEFINE(Primitive_, ([](const py::module *m) {
(void)py::enum_<PrimType>(*m, "prim_type", py::arithmetic())
.value("unknown", PrimType::kPrimTypeUnknown)
.value("builtin", PrimType::kPrimTypeBuiltIn)
.value("py_infer_shape", PrimType::kPrimTypePyInfer)
.value("user_custom", PrimType::kPrimTypeUserCustom)
.value("py_infer_check", PrimType::kPrimTypePyCheck);
(void)py::class_<PrimitivePyAdapter, std::shared_ptr<PrimitivePyAdapter>>(*m, "Primitive_")
.def_readonly(PYTHON_PRIMITIVE_FLAG, &PrimitivePyAdapter::parse_info_)
.def(py::init<py::str &>())
.def("add_attr", &PrimitivePyAdapter::AddPyAttr, "add primitive attr")
.def("del_attr", &PrimitivePyAdapter::DelPyAttr, "del primitive attr")
.def("get_attr_dict", &PrimitivePyAdapter::GetAttrDict, "get primitive attr")
.def("set_prim_type", &PrimitivePyAdapter::set_prim_type, "Set primitive type.")
.def("set_const_prim", &PrimitivePyAdapter::set_const_prim, "Set primitive is const.")
.def("set_const_input_indexes", &PrimitivePyAdapter::set_const_input_indexes,
"Set primitive const input indexes.")
.def("set_signatures", &PrimitivePyAdapter::set_signatures,
"Set primitive inputs signature.")
.def("register_hook", &PrimitivePyAdapter::set_hook, "Set primitive hook function.")
.def("set_instance_name", &PrimitivePyAdapter::set_instance_name,
"Set primitive instance name.");
}));
REGISTER_PYBIND_DEFINE(
Primitive_, ([](const py::module *m) {
(void)py::enum_<PrimType>(*m, "prim_type", py::arithmetic())
.value("unknown", PrimType::kPrimTypeUnknown)
.value("builtin", PrimType::kPrimTypeBuiltIn)
.value("py_infer_shape", PrimType::kPrimTypePyInfer)
.value("user_custom", PrimType::kPrimTypeUserCustom)
.value("py_infer_check", PrimType::kPrimTypePyCheck);
(void)py::class_<PrimitivePyAdapter, std::shared_ptr<PrimitivePyAdapter>>(*m, "Primitive_")
.def_readonly(PYTHON_PRIMITIVE_FLAG, &PrimitivePyAdapter::parse_info_)
.def(py::init<py::str &>())
.def("add_attr", &PrimitivePyAdapter::AddPyAttr, "add primitive attr")
.def("del_attr", &PrimitivePyAdapter::DelPyAttr, "del primitive attr")
.def("get_attr_dict", &PrimitivePyAdapter::GetAttrDict, "get primitive attr")
.def("set_prim_type", &PrimitivePyAdapter::set_prim_type, "Set primitive type.")
.def("set_const_prim", &PrimitivePyAdapter::set_const_prim, "Set primitive is const.")
.def("set_const_input_indexes", &PrimitivePyAdapter::set_const_input_indexes,
"Set primitive const input indexes.")
.def("set_signatures", &PrimitivePyAdapter::set_signatures, "Set primitive inputs signature.")
.def("add_backward_hook_fn", &PrimitivePyAdapter::AddBackwardHookFn, "Add primitive backward hook function.")
.def("remove_backward_hook_fn", &PrimitivePyAdapter::RemoveBackwardHookFn,
"Remove primitive backward hook function.")
.def("set_instance_name", &PrimitivePyAdapter::set_instance_name, "Set primitive instance name.");
}));
} // namespace mindspore

43
mindspore/ccsrc/pybind_api/ir/primitive_py.h
View File

@ -48,17 +48,15 @@ class PrimitivePy : public Primitive {
PrimitivePy(const py::object &python_obj, const PrimitivePyAdapterPtr &adapter);
~PrimitivePy() override;
MS_DECLARE_PARENT(PrimitivePy, Primitive);
const bool parse_info_ = true;
py::function GetBpropFunction();
void set_signatures(const std::vector<Signature> &signatures);
const std::vector<Signature> &signatures() const { return signatures_; }
void CopyHookFunction(const PrimitivePtr &primitive);
py::dict GetAttrDict();
void set_hook(const py::function &hook) { hook_ = hook; }
py::function hook() const { return hook_; }
const std::map<int, py::function> &backward_hook_fn() const { return backward_hook_fn_; }
void CopyHookFunction(const PrimitivePyPtr &primitive_py);
void AddBpropCutPrim(const PrimitivePyPtr &bprop_cut_prim);
void AddBackwardHookFn(const int &key, const py::function &backward_hook_fn);
void RemoveBackwardHookFn(const int &key);
BaseRef RunHookFunction(const VectorRef &args) const;
BaseRef RunCellBpropFunction(const py::tuple &py_args) const;
BaseRef RunCellHookFunction(const py::tuple &py_args) const;
@ -66,12 +64,12 @@ class PrimitivePy : public Primitive {
BaseRef RunComputeFunction(const VectorRef &args) const override;
py::object RunPyComputeFunction(const py::tuple &py_args) const;
bool HasComputeFunction() const;
const bool parse_info_ = true;
py::dict GetAttrDict();
const py::object &GetPyObj() const { return python_obj_; }
py::dict RunInfer(const py::tuple &args);
void RunCheck(const py::tuple &args);
py::object RunInferValue(const py::tuple &args);
bool HasPyObj() const { return python_obj_.operator bool(); }
void RunCheck(const py::tuple &args);
py::dict RunInfer(const py::tuple &args);
py::object RunInferValue(const py::tuple &args);
PrimitivePtr Clone() override;
PrimitivePyAdapterPtr adapter() const { return adapter_; }
void set_bprop_cls_name(const std::string &name) { bprop_cls_name_ = name; }
@ -79,13 +77,14 @@ class PrimitivePy : public Primitive {
private:
py::function GetComputeFunction() const;
void ConvertCTensorToPyTensor(const py::tuple &input_args, py::tuple *convert_args) const;
void CheckHookConsistency(const py::object &grad_out, const py::object &expected_grad_out) const;
void CheckHookConsistency(const py::object &grad_out, const py::object &expected_grad_out, const py::object &code_obj,
const py::object &co_name) const;
py::object python_obj_;
PrimitivePyAdapterPtr adapter_;
py::function hook_;
std::string bprop_cls_name_;
PrimitivePyAdapterPtr adapter_;
std::vector<Signature> signatures_;
std::vector<PrimitivePyWeakPtr> bprop_cut_prims_;
std::map<int, py::function> backward_hook_fn_;
static std::map<std::string, py::object> hook_grad_;
};
@ -96,11 +95,12 @@ class PrimitivePyAdapter {
void AddPyAttr(const py::str &name, const py::object &obj);
void DelPyAttr(const py::str &name);
py::dict GetAttrDict();
int AddBackwardHookFn(const py::function &backward_hook_fn);
void RemoveBackwardHookFn(int key);
void set_prim_type(const PrimType t);
void set_const_prim(bool is_const_prim);
void set_const_input_indexes(const std::vector<size_t> &const_input_indexes);
void set_signatures(const std::vector<Signature> &signatures);
void set_hook(const py::function &hook);
void set_instance_name(const std::string &s);
void set_attached_primitive(const PrimitivePyPtr &prim);
PrimitivePyPtr attached_primitive() const { return attached_primitive_.lock(); }
@ -110,15 +110,16 @@ class PrimitivePyAdapter {
private:
friend PrimitivePy;
bool is_const_prim_{false};
int backward_hook_fn_key_{-1};
std::string name_;
std::string instance_name_;
PrimType prim_type_{kPrimTypeBuiltIn};
PrimitivePyWeakPtr attached_primitive_;
mindspore::HashMap<std::string, ValuePtr> attrs_;
PrimType prim_type_{kPrimTypeBuiltIn};
bool is_const_prim_{false};
std::vector<size_t> const_input_indexes_;
std::vector<Signature> signatures_;
py::function hook_;
std::string instance_name_;
std::map<int, py::function> backward_hook_fn_;
};
} // namespace mindspore
#endif // MINDSPORE_CCSRC_UTILS_PRIMITIVE_PY_H_

1
mindspore/core/base/core_ops.h
View File

@ -768,6 +768,7 @@ GVAR_DEF(PrimitivePtr, kPrimUpdateState, std::make_shared<Primitive>("UpdateStat
GVAR_DEF(PrimitivePtr, kPrimPartial, std::make_shared<Primitive>("Partial", kSideEffectPropagate));
GVAR_DEF(PrimitivePtr, kPrimIdentity, std::make_shared<Primitive>("identity", kSideEffectPropagate));
GVAR_DEF(PrimitivePtr, kPrimHookBackward, std::make_shared<Primitive>("HookBackward"));
GVAR_DEF(PrimitivePtr, kPrimCellBackwardHook, std::make_shared<Primitive>("CellBackwardHook"));
GVAR_DEF(PrimitivePtr, kPrimPrintShapeType, std::make_shared<Primitive>("PrintShapeType"));
GVAR_DEF(PrimitivePtr, kPrimSameTypeShape, std::make_shared<Primitive>("SameTypeShape"));
GVAR_DEF(PrimitivePtr, kPrimPrint, std::make_shared<Primitive>("Print"));

4
mindspore/python/mindspore/_extends/parse/parser.py
View File

@ -123,8 +123,8 @@ def get_parse_method_of_class(obj, parse_method=None):
if parse_method is not None:
method_name = parse_method
elif isinstance(obj, nn.Cell):
if obj.enable_hook:
method_name = "_hook_construct"
if obj.enable_backward_hook:
method_name = "run_backward_hook"
else:
method_name = "construct"
if method_name is not None:

3
mindspore/python/mindspore/common/api.py
View File

@ -505,6 +505,9 @@ class _PynativeExecutor:
def is_top_cell(self):
return self._executor.is_top_cell()
def set_hook_changed(self, cell):
self._executor.set_hook_changed(cell)
def __call__(self, obj, *args, **kwargs):
args = args + tuple(kwargs.values())
return self._executor(obj, args)

110
mindspore/python/mindspore/common/hook_handle.py Normal file
View File

@ -0,0 +1,110 @@
# Copyright 2022 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================
"""The removable handle for cell hook function."""
import weakref
from .api import _pynative_executor
class HookHandle:
r"""
It is the return object of Cell forward pre hook function, forward hook function and backward hook function.
It corresponds to the cell hook function and is used to remove the cell hook function by calling 'remove()'.
Note:
It is only supported in pynative mode and works when registering or removing hook function for Cell object.
Args:
hook_cell (Cell): The Cell object with hook function registered on. Default value: None.
hook_key (int): The key of cell hook function in dict. It is generated during cell hook function registration.
Default value: -1.
hook_type (str): The type of cell hook function: 'forward_pre_hook', 'forward_hook' or 'cell_backward_hook'.
Default value: "".
Returns:
None.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
"""
def __init__(self, hook_cell=None, hook_key=-1, hook_type=""):
if hook_cell is not None:
self._hook_cell = weakref.ref(hook_cell)
else:
self._hook_cell = hook_cell
self._hook_key = hook_key
self._hook_type = hook_type
def remove(self):
"""
Remove the cell hook function, which corresponds to this 'HookHandle' object.
In order to prevent running failed when switching to graph mode, it is not recommended to write in the
construct.
Args:
None.
Returns:
None.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import numpy as np
>>> import mindspore
>>> import mindspore.nn as nn
>>> from mindspore import Tensor
>>> from mindspore import context
>>> from mindspore.ops import GradOperation
>>> context.set_context(mode=context.PYNATIVE_MODE)
>>> def forward_pre_hook_fn(cell_id, inputs):
... print("forward inputs: ", inputs)
...
>>> class Net(nn.Cell):
... def __init__(self):
... super(Net, self).__init__()
... self.mul = nn.MatMul()
... self.handle = self.mul.register_forward_pre_hook(forward_pre_hook_fn)
...
... def construct(self, x, y):
... x = x + x
... x = self.mul(x, y)
... return x
>>> grad = GradOperation(get_all=True)
>>> net = Net()
>>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)), Tensor(np.ones([1]).astype(np.float32)))
forward inputs: (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1],
dtype=Float32, value= [ 1.00000000e+00]))
>>> net.handle.remove()
>>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)), Tensor(np.ones([1]).astype(np.float32)))
>>> print(output)
(Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32,
value= [ 2.00000000e+00]))
"""
if self._hook_cell is not None:
hook_cell = self._hook_cell()
if self._hook_type == "forward_pre_hook" and self._hook_key in hook_cell.forward_pre_hook:
del hook_cell.forward_pre_hook[self._hook_key]
_pynative_executor.set_hook_changed(hook_cell)
elif self._hook_type == "forward_hook" and self._hook_key in hook_cell.forward_hook:
del hook_cell.forward_hook[self._hook_key]
_pynative_executor.set_hook_changed(hook_cell)
elif self._hook_type == "cell_backward_hook":
hook_cell.cell_backward_hook.remove_backward_hook(self._hook_key)
def __del__(self):
self._hook_cell = None
self._hook_key = None
self._hook_type = None

338
mindspore/python/mindspore/nn/cell.py
View File

@ -18,12 +18,13 @@ import inspect
import os
import time
from collections import OrderedDict
from types import FunctionType, MethodType
import numpy
from mindspore._checkparam import args_type_check
from mindspore import log as logger
from mindspore.common.parameter import PARAMETER_NAME_DEFAULT
from mindspore.common.hook_handle import HookHandle
from mindspore.context import ParallelMode
from .. import context
from .._c_expression import init_pipeline, update_func_graph_hyper_params, Cell_, FuncGraph, MixedPrecisionType
@ -32,8 +33,9 @@ from ..common import dtype as mstype
from ..common.api import _cell_graph_executor, _pynative_executor, _check_all_tensor, cells_compile_cache
from ..common.parameter import Parameter, ParameterTuple
from ..common.tensor import Tensor, CSRTensor
from ..ops.operations import HookBackward, Cast
from ..ops.operations import Cast
from ..ops.primitive import Primitive
from ..ops.operations import _inner_ops as inner
from ..parallel._tensor import _load_tensor_by_layout
@ -80,10 +82,10 @@ class Cell(Cell_):
IGNORE_LIST = ['_scope', '_cell_init_args', '_auto_prefix', '_cells', '_params', '_construct_inputs_names',
'_construct_inputs_num', '_create_time', '_mindspore_flags', '_parallel_inputs_run',
'_parameter_layout_dict', '_params_list', '_tensor_list', '_phase',
'_auto_parallel_mode', '_backward_hook', '_bprop_debug', '_is_run', '_param_prefix',
'_attr_synced', 'enable_hook', 'pynative', 'requires_grad',
'_auto_parallel_compile_and_run', 'cell_type']
'_parameter_layout_dict', '_params_list', '_tensor_list', '_phase', '_auto_parallel_mode',
'forward_pre_hook', 'forward_hook', '_enable_forward_pre_hook', '_enable_forward_hook',
'_bprop_debug', 'enable_backward_hook', 'cell_backward_hook', '_is_run', '_param_prefix',
'_attr_synced', 'pynative', 'requires_grad', '_auto_parallel_compile_and_run', 'cell_type']
def __init__(self, auto_prefix=True, flags=None):
Cell_.__init__(self, self._cell_tag)
@ -123,9 +125,13 @@ class Cell(Cell_):
self._parallel_inputs_run = None
if flags:
self.add_flags(**flags)
self._backward_hook = None
self.enable_hook = False
self._bprop_debug = False
self.forward_pre_hook = OrderedDict()
self.forward_hook = OrderedDict()
self._enable_forward_pre_hook = False
self._enable_forward_hook = False
self.enable_backward_hook = False
self.cell_backward_hook = None
self.cell_type = None
self._auto_parallel_compile_and_run = False
self.cast = Cast()
@ -374,10 +380,14 @@ class Cell(Cell_):
self.parameter_broadcast_done = True
def run_construct(self, cast_inputs, kwargs):
if self.enable_hook:
output = self._hook_construct(*cast_inputs)
if self._enable_forward_pre_hook:
cast_inputs = self.run_forward_pre_hook(cast_inputs)
if self.enable_backward_hook:
output = self.run_backward_hook(*cast_inputs)
else:
output = self.construct(*cast_inputs, **kwargs)
if self._enable_forward_hook:
output = self.run_forward_hook(cast_inputs, output)
return output
def _check_construct_args(self, *inputs, **kwargs):
@ -466,9 +476,9 @@ class Cell(Cell_):
# Run in Graph mode.
if context._get_mode() == context.GRAPH_MODE:
self._check_construct_args(*args, **kwargs)
if self.enable_hook:
raise ValueError("For 'Cell', it's not support hook function in graph mode, please use "
"context.set_context to set pynative mode.")
if self._enable_forward_pre_hook or self._enable_forward_hook or self.enable_backward_hook:
logger.warning(f"For 'Cell', it's not support hook function in graph mode, please use "
f"context.set_context to set pynative mode.")
out = self.compile_and_run(*args)
return out
@ -1442,33 +1452,295 @@ class Cell(Cell_):
self.add_flags(auto_parallel=True)
self._get_construct_inputs_number_and_name()
def _hook_construct(self, *inputs):
"""Hook construct method to replace original construct method when hook function enabled."""
inputs = self._backward_hook(*inputs)
inputs = self.construct(inputs)
outputs = self._backward_hook(inputs)
return outputs
def register_backward_hook(self, fn):
def run_forward_pre_hook(self, inputs):
"""
Set the cell backward hook function. Note that this function is only supported in pynative mode.
Note:
fn must be defined as the following code. `cell_name` is the name of registered cell.
`grad_input` is gradient passed to the cell. `grad_output` is the gradient computed and passed to the
next cell or primitive, which may be modified and returned.
hook_fn(cell_name, grad_input, grad_output) -> Tensor or None.
Running forward pre hook function registered in Cell object.
Args:
fn (function): Specifies the hook function with grad as input.
inputs: The input objects of Cell object.
Returns:
- **outputs** - New input objects or none.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
"""
cell_id = self.cls_name + "(" + str(id(self)) + ")"
for fn in self.forward_pre_hook.values():
ret = fn(cell_id, inputs)
if ret is not None:
if not isinstance(ret, tuple):
inputs = (ret,)
else:
inputs = ret
return inputs
def register_forward_pre_hook(self, hook_fn):
"""
Register forward pre hook function for Cell object. Note that this function is only supported in pynative mode.
Note:
'hook_fn' must be defined as the following code.
`cell_id` is the information of registered Cell object. `inputs` is the forward input objects passed to
the Cell. The 'hook_fn' can modify the forward input objects by returning new forward input objects.
It should have the following signature:
hook_fn(cell_id, inputs) -> new input objects or none.
In order to prevent running failed when switching to graph mode, it is not recommended to write in the
construct.
Args:
hook_fn (function): Python function. Forward pre hook function.
Returns:
- **handle** - The handle corresponding to the 'hook_fn'.
Raises:
TypeError: If the `hook_fn` is not a function of python.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import numpy as np
>>> import mindspore
>>> import mindspore.nn as nn
>>> from mindspore import Tensor
>>> from mindspore import context
>>> from mindspore.ops import GradOperation
>>> context.set_context(mode=context.PYNATIVE_MODE)
>>> def forward_pre_hook_fn(cell_id, inputs):
... print("forward inputs: ", inputs)
...
>>> class Net(nn.Cell):
... def __init__(self):
... super(Net, self).__init__()
... self.mul = nn.MatMul()
... self.handle = self.mul.register_forward_pre_hook(forward_pre_hook_fn)
...
... def construct(self, x, y):
... x = x + x
... x = self.mul(x, y)
... return x
>>> grad = GradOperation(get_all=True)
>>> net = Net()
>>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)), Tensor(np.ones([1]).astype(np.float32)))
forward inputs: (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1],
dtype=Float32, value= [ 1.00000000e+00]))
>>> print(output)
(Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32,
value= [ 2.00000000e+00]))
"""
if context.get_context("mode") != context.PYNATIVE_MODE:
logger.warning("Hook function is only supported in pynative mode, you can use context.set_context to set "
"pynative mode.")
logger.warning(f"'register_forward_pre_hook' function is only supported in pynative mode, you can use "
f"context.set_context to set pynative mode.")
return HookHandle()
if not isinstance(hook_fn, (FunctionType, MethodType)):
raise TypeError(f"When using 'register_forward_pre_hook(hook_fn)', the type of 'hook_fn' should be python "
f"function, but got {type(hook_fn)}.")
self._enable_forward_pre_hook = True
_pynative_executor.set_hook_changed(self)
if not hasattr(self, '_forward_pre_hook_key'):
self._forward_pre_hook_key = -1
self._forward_pre_hook_key += 1
self.forward_pre_hook[self._forward_pre_hook_key] = hook_fn
handle = HookHandle(self, self._forward_pre_hook_key, "forward_pre_hook")
return handle
def run_forward_hook(self, inputs, output):
"""
Running forward hook function registered in Cell object.
Args:
inputs: The input objects of Cell object.
output: The output object of Cell object.
Returns:
- **output** - New output object or none.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
"""
cell_id = self.cls_name + "(" + str(id(self)) + ")"
for fn in self.forward_hook.values():
ret = fn(cell_id, inputs, output)
if ret is not None:
output = ret
return output
def register_forward_hook(self, hook_fn):
"""
Set the cell forward hook function. Note that this function is only supported in pynative mode.
Note:
'hook_fn' must be defined as the following code.
`cell_id` is the information of registered Cell object. `inputs` is the forward input objects passed to
the Cell. `output` is the forward output object of the Cell. The 'hook_fn' can modify the forward output
object by returning new forward output object.
It should have the following signature:
hook_fn(cell_id, inputs, output) -> new output object or none.
In order to prevent running failed when switching to graph mode, it is not recommended to write in the
construct.
Args:
hook_fn (function): Python function. Forward hook function.
Returns:
- **handle** - The handle corresponding to the 'hook_fn'.
Raises:
TypeError: If the `hook_fn` is not a function of python.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import numpy as np
>>> import mindspore
>>> import mindspore.nn as nn
>>> from mindspore import Tensor
>>> from mindspore import context
>>> from mindspore.ops import GradOperation
>>> context.set_context(mode=context.PYNATIVE_MODE)
>>> def forward_hook_fn(cell_id, inputs, output):
... print("forward inputs: ", inputs)
... print("forward output: ", output)
...
>>> class Net(nn.Cell):
... def __init__(self):
... super(Net, self).__init__()
... self.mul = nn.MatMul()
... self.handle = self.mul.register_forward_hook(forward_hook_fn)
...
... def construct(self, x, y):
... x = x + x
... x = self.mul(x, y)
... return x
>>> grad = GradOperation(get_all=True)
>>> net = Net()
>>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)), Tensor(np.ones([1]).astype(np.float32)))
forward inputs: (Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1],
dtype=Float32, value= [ 1.00000000e+00]))
forward output: 2.0
>>> print(output)
(Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]), Tensor(shape=[1], dtype=Float32,
value= [ 2.00000000e+00]))
"""
if context.get_context("mode") != context.PYNATIVE_MODE:
logger.warning(f"'register_forward_hook' function is only supported in pynative mode, you can use "
f"context.set_context to set pynative mode.")
return HookHandle()
if not isinstance(hook_fn, (FunctionType, MethodType)):
raise TypeError(f"When using 'register_forward_hook(hook_fn)', the type of 'hook_fn' should be python "
f"function, but got {type(hook_fn)}.")
self._enable_forward_hook = True
_pynative_executor.set_hook_changed(self)
if not hasattr(self, '_forward_hook_key'):
self._forward_hook_key = -1
self._forward_hook_key += 1
self.forward_hook[self._forward_hook_key] = hook_fn
handle = HookHandle(self, self._forward_hook_key, "forward_hook")
return handle
def run_backward_hook(self, *inputs):
"""
Backward hook construct method to replace original construct method.
Args:
inputs: The input objects of Cell object.
Returns:
- **outputs** - The output objects of Cell object.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
"""
inputs = self.cell_backward_hook(*inputs)
if isinstance(inputs, tuple):
inputs = self.construct(*inputs)
else:
self._backward_hook = HookBackward(fn, self.cls_name + "(" + str(id(self)) + ")")
self.enable_hook = True
inputs = self.construct(inputs)
if isinstance(inputs, tuple):
outputs = self.cell_backward_hook(*inputs)
else:
outputs = self.cell_backward_hook(inputs)
return outputs
def register_backward_hook(self, hook_fn):
"""
Register the backward hook function. Note that this function is only supported in pynative mode.
Note:
The 'hook_fn' must be defined as the following code.
`cell_id` is the information of registered cell. `grad_input` is the gradient passed to the cell.
`grad_output` is the gradient computed and passed to the next cell or primitive, which may be modified by
returning a new output gradient.
The 'hook_fn' should have the following signature:
hook_fn(cell_id, grad_input, grad_output) -> New output gradient or none.
The 'hook_fn' is executed in the python environment.
In order to prevent running failed when switching to graph mode, it is not recommended to write in the
construct.
Args:
hook_fn (function): Python function. Backward hook function.
Returns:
- **handle** - The handle corresponding to the 'hook_fn'.
Raises:
TypeError: If the `hook_fn` is not a function of python.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import numpy as np
>>> import mindspore
>>> import mindspore.nn as nn
>>> from mindspore import Tensor
>>> from mindspore import context
>>> from mindspore.ops import GradOperation
>>> context.set_context(mode=context.PYNATIVE_MODE)
>>> def backward_hook_fn(cell_id, grad_input, grad_output):
... print("backward input: ", grad_input)
... print("backward output: ", grad_output)
...
>>> class Net(nn.Cell):
... def __init__(self):
... super(Net, self).__init__()
... self.relu = nn.ReLU()
... self.handle = self.relu.register_backward_hook(backward_hook_fn)
...
... def construct(self, x):
... x = x + x
... x = self.relu(x)
... return x
>>> grad = GradOperation(get_all=True)
>>> net = Net()
>>> output = grad(net)(Tensor(np.ones([1]).astype(np.float32)))
backward input: (Tensor(shape=[1], dtype=Float32, value= [ 1.00000000e+00]),)
backward output: (Tensor(shape=[1], dtype=Float32, value= [ 1.00000000e+00]),)
>>> print(output)
(Tensor(shape=[1], dtype=Float32, value= [ 2.00000000e+00]),)
"""
if context.get_context("mode") != context.PYNATIVE_MODE:
logger.warning(f"'register_backward_hook' function is only supported in pynative mode, you can use "
f"context.set_context to set pynative mode.")
return HookHandle()
if not isinstance(hook_fn, (FunctionType, MethodType)):
raise TypeError(f"When using 'register_backward_hook(hook_fn)', the type of 'hook_fn' should be python "
f"function, but got {type(hook_fn)}.")
if self.cell_backward_hook is None:
self.enable_backward_hook = True
self.cell_backward_hook = inner.CellBackwardHook(self.cls_name + "(" + str(id(self)) + ")")
backward_hook_key = self.cell_backward_hook.register_backward_hook(hook_fn)
handle = HookHandle(self, backward_hook_key, "cell_backward_hook")
else:
backward_hook_key = self.cell_backward_hook.register_backward_hook(hook_fn)
handle = HookHandle(self, backward_hook_key, "cell_backward_hook")
return handle
def set_param_ps(self, recurse=True, init_in_server=False):
"""

118
mindspore/python/mindspore/ops/operations/_inner_ops.py
View File

@ -14,7 +14,7 @@
# ============================================================================
"""Inner operators."""
from types import FunctionType, MethodType
import numpy as np
from mindspore.common import Tensor
@ -1702,3 +1702,119 @@ class ParallelResizeBilinear(PrimitiveWithInfer):
return {'shape': output_shape,
'dtype': x_dtype,
'value': None}
class CellBackwardHook(PrimitiveWithInfer):
r"""
This operator is used to hook input gradient and output gradient of Cell object.
Note:
This operator is only used in backward hook function of Cell object in pynative mode.
Args:
cell_id (str): Used to identify which cell obj the hook function registered on. For example, 'nn.Add()' is a
cell object.
Inputs:
- **input** - The variable to hook.
Outputs:
- **output** - Returns `input` directly. `CellBackwardHook` does not affect the forward result.
Supported Platforms:
``Ascend`` ``GPU`` ``CPU``
Examples:
>>> import mindspore
>>> from mindspore import Tensor
>>> from mindspore import context
>>> from mindspore.ops import GradOperation
>>> from mindspore.ops.operations import _inner_ops as inner
>>> context.set_context(mode=context.PYNATIVE_MODE)
>>> def hook_fn(grad):
... print(grad)
...
>>> hook = inner.CellBackwardHook()
>>> hook_fn_key = hook.register_backward_hook(hook_fn)
>>> def hook_test(x, y):
... z = x * y
... z = hook(z)
... z = z * y
... return z
...
>>> grad_all = GradOperation(get_all=True)
>>> def backward(x, y):
... return grad_all(hook_test)(x, y)
...
>>> output = backward(Tensor(1, mindspore.float32), Tensor(2, mindspore.float32))
(Tensor(shape=[], dtype=Float32, value= 2),)
>>> print(output)
(Tensor(shape=[], dtype=Float32, value= 4), Tensor(shape=[], dtype=Float32, value= 4))
>>> hook.remove_backward_hook(hook_fn_key)
>>> output = backward(Tensor(1, mindspore.float32), Tensor(2, mindspore.float32))
>>> print(output)
(Tensor(shape=[], dtype=Float32, value= 4), Tensor(shape=[], dtype=Float32, value= 4))
"""
def __init__(self, cell_id=""):
"""Initialize CellBackwardHook"""
super(CellBackwardHook, self).__init__(self.__class__.__name__)
self.cell_id = cell_id
self.add_prim_attr("cell_id", cell_id)
self.init_attrs["cell_id"] = cell_id
def infer_shape(self, *inputs_shape):
if len(inputs_shape) == 1:
return inputs_shape[0]
return inputs_shape
def infer_dtype(self, *inputs_type):
if len(inputs_type) == 1:
return inputs_type[0]
return inputs_type
def register_backward_hook(self, hook_fn):
r"""
This function is used to register backward hook function. Note that this function is only supported in pynative
mode.
Note:
The 'hook_fn' must be defined as the following code.
`cell_id` is the information of registered cell. `grad_input` is the gradient passed to the cell.
`grad_output` is the gradient computed and passed to the next cell or primitive, which may be modified by
returning a new output gradient.
The 'hook_fn' should have the following signature:
hook_fn(cell_id, grad_input, grad_output) -> New output gradient or none.
The 'hook_fn' is executed in the python environment.
Args:
hook_fn (Function): Python function. Backward hook function.
Returns:
- **key** (int) - The key of 'hook_fn'.
Raises:
TypeError: If the `hook_fn` is not a function of python.
"""
if not isinstance(hook_fn, (FunctionType, MethodType)):
raise TypeError(f"When using 'register_backward_hook(hook_fn)', the type of 'hook_fn' should be python "
f"function, but got {type(hook_fn)}.")
key = self.add_backward_hook_fn(hook_fn)
return key
def remove_backward_hook(self, key):
r"""
This function is used to remove backward hook function. Note that this operation is only supported in pynative
mode.
Note:
The 'key' is the object returned by 'register_backward_hook' function of the same CellBackwardHook
operator.
Args:
key (int): The key corresponding to the 'hook_fn'.
Returns:
None.
"""
self.remove_backward_hook_fn(key)

45
mindspore/python/mindspore/ops/operations/debug_ops.py
View File

@ -324,28 +324,25 @@ class InsertGradientOf(PrimitiveWithInfer):
class HookBackward(PrimitiveWithInfer):
"""
This operation is used as a tag to hook gradient in intermediate variables. Note that this function
is only supported in Pynative Mode.
is only supported in pynative mode.
Note:
The hook function must be defined like `hook_fn(grad) -> Tensor or None`,
where grad is the gradient passed to the primitive and gradient may be
modified and passed to next primitive. The difference between a hook function and
callback of InsertGradientOf is that a hook function is executed in the python
environment while callback will be parsed and added to the graph.
The hook function must be defined like `hook_fn(grad) -> new gradient or None`, where the 'grad' is the
gradient passed to the primitive. The 'grad' may be modified by returning a new gradient and passed to next
primitive. The difference between a hook function and callback of InsertGradientOf is that the hook function is
executed in the python environment while callback will be parsed and added to the graph.
Args:
hook_fn (Function): Python function. hook function.
cell_id (str): Used to identify whether the function registered by the hook is actually registered on
the specified Cell. Where the Cell is an object. For example, 'nn.Add' is a Cell object.
The default value of cell_id is empty string(""), in this case, the system will automatically
register when registering. Add a value of cell_id,
the value of cell_id currently does not support custom values.
Inputs:
- **inputs** (Tensor) - The variable to hook.
- **input** (Tensor) - The variable to hook.
Outputs:
- **output** (Tensor) - Returns `input` directly. `HookBackward` does not affect the forward result.
Raises:
TypeError: If `inputs` are not a Tensor.
TypeError: If `input` is not a tensor.
TypeError: If `hook_fn` is not a function of python.
Supported Platforms:
@ -353,15 +350,14 @@ class HookBackward(PrimitiveWithInfer):
Examples:
>>> import mindspore
>>> from mindspore import context
>>> from mindspore import Tensor
>>> from mindspore import ops
>>> from mindspore import Tensor
>>> from mindspore import context
>>> from mindspore.ops import GradOperation
>>> context.set_context(mode=context.PYNATIVE_MODE, device_target="GPU")
>>> def hook_fn(grad_out):
... print(grad_out)
>>> context.set_context(mode=context.PYNATIVE_MODE)
>>> def hook_fn(grad):
... print(grad)
...
>>> grad_all = GradOperation(get_all=True)
>>> hook = ops.HookBackward(hook_fn)
>>> def hook_test(x, y):
... z = x * y
@ -369,6 +365,7 @@ class HookBackward(PrimitiveWithInfer):
... z = z * y
... return z
...
>>> grad_all = GradOperation(get_all=True)
>>> def backward(x, y):
... return grad_all(hook_test)(x, y)
...
@ -378,16 +375,16 @@ class HookBackward(PrimitiveWithInfer):
(Tensor(shape=[], dtype=Float32, value= 4), Tensor(shape=[], dtype=Float32, value= 4))
"""
def __init__(self, hook_fn, cell_id=""):
def __init__(self, hook_fn):
"""Initialize HookBackward."""
super(HookBackward, self).__init__(self.__class__.__name__)
self.add_prim_attr("cell_id", cell_id)
self.init_attrs["cell_id"] = cell_id
if not isinstance(hook_fn, (FunctionType, MethodType)):
raise TypeError(f"For '{self.name}', the type of 'hook_fn' should be python function, "
f"but got {type(hook_fn)}.")
self.register_hook(hook_fn)
self.cell_id = cell_id
self.add_prim_attr("cell_id", "")
self.init_attrs["cell_id"] = ""
self.cell_id = ""
self.add_backward_hook_fn(hook_fn)
def infer_shape(self, *inputs_shape):
if len(inputs_shape) == 1: