diff --git a/mindspore/_extends/builtin_operations.py b/mindspore/_extends/builtin_operations.py index 6bd382c1b6..04dd443bde 100644 --- a/mindspore/_extends/builtin_operations.py +++ b/mindspore/_extends/builtin_operations.py @@ -14,6 +14,7 @@ # ============================================================================ """builtin_operations""" import numpy as np +from mindspore.ops import functional as F from mindspore.common.tensor import Tensor from mindspore.common.dtype import dtype_to_nptype, get_py_obj_dtype @@ -171,3 +172,13 @@ def tuple_to_array(x): def stop_gradient(x): """Implement `stop_gradient`.""" return x + + +def mixed_precision_cast(dst_type, x): + """Implement `mixed_precision_cast`.""" + if isinstance(x, tuple): + res = list() + for item in x: + res.append(F.cast(item, dst_type)) + return tuple(res) + return F.cast(x, dst_type) diff --git a/mindspore/ccsrc/pipeline/pynative/base.h b/mindspore/ccsrc/pipeline/pynative/base.h index 2a380d6900..eed522e7df 100644 --- a/mindspore/ccsrc/pipeline/pynative/base.h +++ b/mindspore/ccsrc/pipeline/pynative/base.h @@ -59,7 +59,7 @@ struct OpExecInfo { using OpExecInfoPtr = std::shared_ptr; OpExecInfoPtr GenerateOpExecInfo(const py::args &args, py::list *const out_args); -const std::set ignore_infer_prim = {"make_ref"}; +const std::set ignore_infer_prim = {"make_ref", "mixed_precision_cast"}; } // namespace pynative } // namespace mindspore diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc index a39a931b29..98b9e5b62d 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.cc @@ -57,7 +57,7 @@ using mindspore::tensor::TensorPy; const char SINGLE_OP_GRAPH[] = "single_op_graph"; // primitive unable to infer value for constant input in PyNative mode -const std::set vm_operators = {"make_ref", "HookBackward", "stop_gradient"}; +const std::set vm_operators = {"make_ref", "HookBackward", "stop_gradient", "mixed_precision_cast"}; namespace mindspore { namespace pynative { @@ -766,6 +766,9 @@ PynativeExecutor::PynativeExecutor() { grad_flag_ = false; } void PynativeExecutor::NewGraphInner(const py::object &cell, const py::args &args) { auto cell_id = GetId(cell); if (cell_graph_map_.count(cell_id) != 0) { + if (cell_resource_map_.find(cell_id) != cell_resource_map_.end()) { + resource_ = cell_resource_map_[cell_id]; + } MS_LOG(DEBUG) << "Newgraph already compiled"; return; } @@ -774,6 +777,8 @@ void PynativeExecutor::NewGraphInner(const py::object &cell, const py::args &arg if (top_g_ == nullptr) { top_g_ = curr_g_ = g; + resource_ = std::make_shared(); + cell_resource_map_[cell_id] = resource_; df_builder_ = std::make_shared(); MS_LOG(DEBUG) << "First new graph" << top_g_.get(); Pushp(); @@ -1075,6 +1080,7 @@ void PynativeExecutor::Clear(const std::string &flag) { MS_LOG(INFO) << "Clear res"; (void)graph_map_.erase(flag); (void)cell_graph_map_.erase(flag); + (void)cell_resource_map_.erase(flag); Clean(); // Maybe exit in the pynative runing op, so need reset pynative flag. auto ms_context = MsContext::GetInstance(); @@ -1086,6 +1092,7 @@ void PynativeExecutor::Clear(const std::string &flag) { MS_LOG(INFO) << "Clear"; top_g_ = nullptr; + df_builder_ = nullptr; curr_g_ = nullptr; graph_info_map_.clear(); std::stack().swap(graph_p_); @@ -1095,7 +1102,6 @@ void PynativeExecutor::Clean() { MS_LOG(INFO) << "Clean all res"; Clear(); grad_flag_ = false; - df_builder_ = nullptr; ad::CleanRes(); pipeline::ReclaimOptimizer(); } diff --git a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h index 4a2c525534..850dc56691 100644 --- a/mindspore/ccsrc/pipeline/pynative/pynative_execute.h +++ b/mindspore/ccsrc/pipeline/pynative/pynative_execute.h @@ -115,6 +115,7 @@ class PynativeExecutor : public std::enable_shared_from_this { bool grad_flag_; std::unordered_map graph_map_; std::unordered_map cell_graph_map_; + std::unordered_map cell_resource_map_; std::unordered_map graph_info_map_; std::stack graph_p_; FuncGraphPtr top_g_; diff --git a/mindspore/nn/wrap/loss_scale.py b/mindspore/nn/wrap/loss_scale.py index 31dfb9f9ed..9cda88b26b 100644 --- a/mindspore/nn/wrap/loss_scale.py +++ b/mindspore/nn/wrap/loss_scale.py @@ -206,6 +206,7 @@ class TrainOneStepWithLossScaleCell(Cell): def __init__(self, network, optimizer, scale_update_cell=None): super(TrainOneStepWithLossScaleCell, self).__init__(auto_prefix=False) self.network = network + self.network.set_grad() self.network.add_flags(defer_inline=True) self.weights = optimizer.parameters self.optimizer = optimizer diff --git a/tests/ut/python/ops/test_control_ops.py b/tests/ut/python/ops/test_control_ops.py index 53b42b8f66..24c593c48c 100644 --- a/tests/ut/python/ops/test_control_ops.py +++ b/tests/ut/python/ops/test_control_ops.py @@ -20,6 +20,7 @@ import mindspore as ms from mindspore import Tensor from mindspore import context from mindspore import nn +from mindspore.common import dtype as mstype from mindspore.ops import composite as C from mindspore.ops import functional as F from mindspore.ops import operations as P @@ -639,3 +640,9 @@ def test_large_for_loop_with_continue_break(): t = Tensor(np.ones([2, 3], dtype=np.float32)) net = Net() net(t) + + +def test_mixed_precision_cast(): + x = Tensor(np.ones([2, 3], dtype=np.float32)) + z = F.mixed_precision_cast(mstype.float16, x) + assert z.dtype == mstype.float16