Star1371
2131
代码
代码
Issues796
Pull Requests717
Actions
项目讨论
Wiki
分析
Star1371
2131
已合并
fix inplace_apply_keras_momentum 注册图模式接口修复 #8877
Davon创建于 8月18日
fix inplace_apply_keras_momentum 注册图模式接口修复 #8877
已合并
Davon创建于 8月18日
从Davon14272/ops-nn_backupfix/inplace-apply-keras-momentum-infer-dtype合入到cann/ops-nnmaster
共 2 个文件变更+21-16
Roptim/inplace_apply_keras_momentum/op_host/inplace_apply_keras_momentum_inferdatatype.cpp→optim/inplace_apply_keras_momentum/op_graph/inplace_apply_keras_momentum_graph_infer.cpp+5-8
@@ -11,17 +11,14 @@
11/* Generated By CANNBot */11/* Generated By CANNBot */
12 12 
13/*!13/*!
14- * \file inplace_apply_keras_momentum_inferdatatype.cpp14+ * \file inplace_apply_keras_momentum_graph_infer.cpp
15- * \brief InplaceApplyKerasMomentum 数据类型推导：var/accum 双输出 dtype15+ * \brief InplaceApplyKerasMomentum graph dtype inference
16 */16 */
17 17 
18#include "register/op_impl_registry.h"18#include "register/op_impl_registry.h"
19-#include "exe_graph/runtime/inferDataType_context.h"
20-#include "op_common/log/log.h"
21- 
22-using namespace ge;
23 19 
24namespace ops {20namespace ops {
21+using namespace ge;
25 22 
26static constexpr size_t VAR_INDEX = 0;23static constexpr size_t VAR_INDEX = 0;
27static constexpr size_t ACCUM_INDEX = 1;24static constexpr size_t ACCUM_INDEX = 1;
@@ -29,10 +26,10 @@ static constexpr size_t ACCUM_INDEX = 1;
29static graphStatus InferDataType4InplaceApplyKerasMomentum(gert::InferDataTypeContext* context)26static graphStatus InferDataType4InplaceApplyKerasMomentum(gert::InferDataTypeContext* context)
30{27{
31 context->SetOutputDataType(VAR_INDEX, context->GetInputDataType(VAR_INDEX));28 context->SetOutputDataType(VAR_INDEX, context->GetInputDataType(VAR_INDEX));
32- context->SetOutputDataType(ACCUM_INDEX, context->GetInputDataType(VAR_INDEX));29+ context->SetOutputDataType(ACCUM_INDEX, context->GetInputDataType(ACCUM_INDEX));
33 return GRAPH_SUCCESS;30 return GRAPH_SUCCESS;
34}31}
35 32 
36-IMPL_OP_INFERSHAPE(InplaceApplyKerasMomentum).InferDataType(InferDataType4InplaceApplyKerasMomentum);33+IMPL_OP(InplaceApplyKerasMomentum).InferDataType(InferDataType4InplaceApplyKerasMomentum);
37 34 
38} // namespace ops35} // namespace ops
Moptim/inplace_apply_keras_momentum/tests/ut/op_host/test_inplace_apply_keras_momentum_infershape.cpp+16-8
@@ -21,7 +21,7 @@
21 * - output var (index 0).shape = input var (index 0).shape21 * - output var (index 0).shape = input var (index 0).shape
22 * - output accum (index 1).shape = input var (index 0).shape22 * - output accum (index 1).shape = input var (index 0).shape
23 * - output var.dtype = input var.dtype23 * - output var.dtype = input var.dtype
24- * - output accum.dtype = input var.dtype24+ * - output accum.dtype = input accum.dtype
25 */25 */
26 26 
27#include <cstdint>27#include <cstdint>
@@ -104,13 +104,20 @@ static std::pair<Shape, Shape> InferShape4InplaceApplyKerasMomentum(const Shape&
104 *104 *
105 * 原始逻辑：105 * 原始逻辑：
106 * context->SetOutputDataType(VAR_INDEX, context->GetInputDataType(VAR_INDEX));106 * context->SetOutputDataType(VAR_INDEX, context->GetInputDataType(VAR_INDEX));
107- * context->SetOutputDataType(ACCUM_INDEX, context->GetInputDataType(VAR_INDEX));107+ * context->SetOutputDataType(ACCUM_INDEX, context->GetInputDataType(ACCUM_INDEX));
108 *108 *
109- * 双输出 dtype 推导：两个输出 dtype 均与输入 var 相同109+ * lr/momentum 固定为 FP32，不参与输出 dtype 推导。
110 */110 */
111-static std::pair<LocalDataType, LocalDataType> InferDataType4InplaceApplyKerasMomentum(LocalDataType inputVarDtype)111+static std::pair<LocalDataType, LocalDataType> InferDataType4InplaceApplyKerasMomentum(LocalDataType inputVarDtype,
112+ LocalDataType inputAccumDtype,
113+ LocalDataType inputLrDtype,
114+ LocalDataType inputGradDtype,
115+ LocalDataType inputMomentumDtype)
112{116{
113- return {inputVarDtype, inputVarDtype};117+ (void)inputLrDtype;
118+ (void)inputGradDtype;
119+ (void)inputMomentumDtype;
120+ return {inputVarDtype, inputAccumDtype};
114}121}
115 122 
116// ===================== 测试用例 =====================123// ===================== 测试用例 =====================
@@ -238,7 +245,7 @@ void test_infershape_unknown_shape()
238void test_infer_dtype_fp32()245void test_infer_dtype_fp32()
239{246{
240 TEST_BEGIN("infer_dtype_fp32");247 TEST_BEGIN("infer_dtype_fp32");
241- auto [dtVar, dtAccum] = InferDataType4InplaceApplyKerasMomentum(DT_FLOAT);248+ auto [dtVar, dtAccum] = InferDataType4InplaceApplyKerasMomentum(DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT);
242 ASSERT_EQ(dtVar, DT_FLOAT);249 ASSERT_EQ(dtVar, DT_FLOAT);
243 ASSERT_EQ(dtAccum, DT_FLOAT);250 ASSERT_EQ(dtAccum, DT_FLOAT);
244 TEST_END();251 TEST_END();
@@ -248,7 +255,8 @@ void test_infer_dtype_fp32()
248void test_infer_dtype_fp16()255void test_infer_dtype_fp16()
249{256{
250 TEST_BEGIN("infer_dtype_fp16");257 TEST_BEGIN("infer_dtype_fp16");
251- auto [dtVar, dtAccum] = InferDataType4InplaceApplyKerasMomentum(DT_FLOAT16);258+ auto [dtVar, dtAccum] = InferDataType4InplaceApplyKerasMomentum(DT_FLOAT16, DT_FLOAT16, DT_FLOAT, DT_FLOAT16,
259+ DT_FLOAT);
252 ASSERT_EQ(dtVar, DT_FLOAT16);260 ASSERT_EQ(dtVar, DT_FLOAT16);
253 ASSERT_EQ(dtAccum, DT_FLOAT16);261 ASSERT_EQ(dtAccum, DT_FLOAT16);
254 TEST_END();262 TEST_END();
@@ -258,7 +266,7 @@ void test_infer_dtype_fp16()
258void test_infer_dtype_bf16()266void test_infer_dtype_bf16()
259{267{
260 TEST_BEGIN("infer_dtype_bf16");268 TEST_BEGIN("infer_dtype_bf16");
261- auto [dtVar, dtAccum] = InferDataType4InplaceApplyKerasMomentum(DT_BF16);269+ auto [dtVar, dtAccum] = InferDataType4InplaceApplyKerasMomentum(DT_BF16, DT_BF16, DT_FLOAT, DT_BF16, DT_FLOAT);
262 ASSERT_EQ(dtVar, DT_BF16);270 ASSERT_EQ(dtVar, DT_BF16);
263 ASSERT_EQ(dtAccum, DT_BF16);271 ASSERT_EQ(dtAccum, DT_BF16);
264 TEST_END();272 TEST_END();