p47068915
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!24671 Enable CPU BatchNorm grad in eval mode 

Merge pull request !24671 from huangbo/master_codedex
This commit is contained in:
i-robot 2021-10-12 02:09:00 +00:00 committed by Gitee
parent 6250998108 70ebd9bdbc
commit 4caa16dee8
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mindspore/ccsrc/backend/kernel_compiler/cpu/mkldnn/batch_norm_grad_cpu_kernel.cc
View File

@ -56,8 +56,13 @@ void BatchNormGradCPUKernel::InitKernel(const CNodePtr &kernel_node) {
dnnl::memory::desc x_desc = GetDefaultMemDesc(x_shape);
dnnl::memory::desc scale_bias_desc = GetDefaultMemDesc({SCALE_SHIFT_NUM, channel});
auto epsilon = AnfAlgo::GetNodeAttr<float>(kernel_node, "epsilon");
auto prop_kind = dnnl::prop_kind::forward_training;
auto normalization_flags = dnnl::normalization_flags::use_scale_shift;
auto is_train = AnfAlgo::GetNodeAttr<bool>(kernel_node, "is_training");
auto prop_kind = dnnl::prop_kind::forward_inference;
auto normalization_flags = dnnl::normalization_flags::use_scale_shift | dnnl::normalization_flags::use_global_stats;
if (is_train) {
prop_kind = dnnl::prop_kind::forward_training;
normalization_flags = dnnl::normalization_flags::use_scale_shift;
}
// fused Batch Normalization forward description
dnnl::batch_normalization_forward::desc desc =