Skip to content

optimization batch_norm 2D and NCHW format on CPU #34585

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Aug 9, 2021
Merged

optimization batch_norm 2D and NCHW format on CPU #34585

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b881367
optimization channel
Zjq9409 Aug 3, 2021
c51c122
support NCHW multi-threading
Zjq9409 Aug 4, 2021
60cc4eb
test
Zjq9409 Aug 5, 2021
37c3599
add multi-threading
Zjq9409 Aug 5, 2021
0f3c528
optimization tow dim NCHW format
Zjq9409 Aug 6, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 14 additions & 4 deletions paddle/fluid/operators/batch_norm_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -295,8 +295,7 @@ class BatchNormKernel<platform::CPUDeviceContext, T>
bool global_stats = test_mode || use_global_stats;

const std::string data_layout_str = ctx.Attr<std::string>("data_layout");
const DataLayout data_layout =
framework::StringToDataLayout(data_layout_str);
DataLayout data_layout = framework::StringToDataLayout(data_layout_str);

const auto *x = ctx.Input<Tensor>("X");
const auto &x_dims = x->dims();
Expand Down Expand Up @@ -332,6 +331,12 @@ class BatchNormKernel<platform::CPUDeviceContext, T>
saved_mean->mutable_data<T>(ctx.GetPlace());
saved_variance->mutable_data<T>(ctx.GetPlace());

// input dimension is 2 and the format is NCHW. The input can be regarded
// as NHWC format
if (x_dims.size() == 2 && data_layout == DataLayout::kNCHW) {
Copy link
Contributor

@Xreki Xreki Aug 9, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

判断条件其实可以简化下，是不是if (sample_size == 1)就可以了？以及代码书写也可以简化下：const DataLayout data_layout = sample_size == 1 ? DataLayout::kNHWC : framework::StringToDataLayout(data_layout_str);

Copy link
Contributor Author

@Zjq9409 Zjq9409 Aug 9, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

经过验证，在sample_size为1时，在代码"const int C = (data_layout == DataLayout::kNCHW ? x_dims[1]: x_dims[x_dims.size() - 1]);" 执行完成后获取到C的值，再使用sample_size==1判断修改layout是可以的。根据建议修改

Copy link
Contributor Author

@Zjq9409 Zjq9409 Aug 9, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

但是对于简化的写法应该是不可行的，因为如果这里直接修改data_layout的话，后续会通过这个data_layout来计算C的值，如果输入为[10,10,1,1]，输入格式为"NCHW"，sample_size==1，满足条件，这时输入格式改为"NHWC"，后续计算C会等于1，而实际上C=10.

data_layout = DataLayout::kNHWC;
}

if (!global_stats) {
// saved_xx is use just in this batch of data
EigenVectorArrayMap<T> saved_mean_e(
Expand Down Expand Up @@ -578,8 +583,7 @@ class BatchNormGradKernel<platform::CPUDeviceContext, T>
bool use_global_stats = ctx.Attr<bool>("use_global_stats");
const bool is_test = ctx.Attr<bool>("is_test");
const float epsilon = ctx.Attr<float>("epsilon");
const DataLayout data_layout =
framework::StringToDataLayout(data_layout_str);
DataLayout data_layout = framework::StringToDataLayout(data_layout_str);

auto *d_x = ctx.Output<Tensor>(framework::GradVarName("X"));
auto *d_scale = ctx.Output<Tensor>(framework::GradVarName("Scale"));
Expand Down Expand Up @@ -633,6 +637,12 @@ class BatchNormGradKernel<platform::CPUDeviceContext, T>
: x_dims[x_dims.size() - 1]);
const int sample_size = x->numel() / N / C;

// input dimension is 2 and the format is NCHW. The input can be regarded as
// NHWC format
if (x_dims.size() == 2 && data_layout == DataLayout::kNCHW) {
data_layout = DataLayout::kNHWC;
}

// init output
if (d_x) {
d_x->mutable_data<T>(ctx.GetPlace());
Expand Down