batch_norm_op.cc 25.0 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
Q
Qiao Longfei 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14

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. */

Y
Yi Wang 已提交
15
#include "paddle/fluid/operators/batch_norm_op.h"
Q
qingqing01 已提交
16
#include <memory>
S
Siddharth Goyal 已提交
17
#include <string>
Q
qingqing01 已提交
18
#include <unordered_map>
Y
Yi Wang 已提交
19
#include "paddle/fluid/framework/data_layout.h"
20 21 22
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif
Q
Qiao Longfei 已提交
23 24 25 26

namespace paddle {
namespace operators {

Q
qingqing01 已提交
27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48
void BatchNormOp::InferShape(framework::InferShapeContext *ctx) const {
  PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) of ConvOp should not be null.");
  PADDLE_ENFORCE(ctx->HasInput("Scale"),
                 "Input(Scale) of ConvOp should not be null.");
  PADDLE_ENFORCE(ctx->HasInput("Bias"),
                 "Input(Bias) of ConvOp should not be null.");
  PADDLE_ENFORCE(ctx->HasInput("Mean"),
                 "Input(Mean) of ConvOp should not be null.");
  PADDLE_ENFORCE(ctx->HasInput("Variance"),
                 "Input(Variance) of ConvOp should not be null.");
  PADDLE_ENFORCE(ctx->HasOutput("Y"),
                 "Output(Y) of ConvOp should not be null.");
  bool is_test = ctx->Attrs().Get<bool>("is_test");
  if (!is_test) {
    PADDLE_ENFORCE(ctx->HasOutput("MeanOut"),
                   "Output(MeanOut) of ConvOp should not be null.");
    PADDLE_ENFORCE(ctx->HasOutput("VarianceOut"),
                   "Output(VarianceOut) of ConvOp should not be null.");
    PADDLE_ENFORCE(ctx->HasOutput("SavedMean"),
                   "Output(SavedMean) of ConvOp should not be null.");
    PADDLE_ENFORCE(ctx->HasOutput("SavedVariance"),
                   "Output(SavedVariance) of ConvOp should not be null.");
Q
Qiao Longfei 已提交
49
  }
K
Kexin Zhao 已提交
50

Q
qingqing01 已提交
51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
  // make sure Mean/MeanOut and Variance/VarianceOut share memory in Python
  PADDLE_ENFORCE_EQ(ctx->Inputs("Mean")[0], ctx->Outputs("MeanOut")[0],
                    "Mean and MeanOut should share the same memory");
  PADDLE_ENFORCE_EQ(ctx->Inputs("Variance")[0], ctx->Outputs("VarianceOut")[0],
                    "Variance and VarianceOut should share the same memory");

  const auto x_dims = ctx->GetInputDim("X");
  const DataLayout data_layout = framework::StringToDataLayout(
      ctx->Attrs().Get<std::string>("data_layout"));

  PADDLE_ENFORCE(x_dims.size() >= 2 && x_dims.size() <= 5,
                 "Input X must have 2 to 5 dimensions.");

  const int64_t C =
      (data_layout == DataLayout::kNCHW ? x_dims[1]
                                        : x_dims[x_dims.size() - 1]);

  PADDLE_ENFORCE_EQ(ctx->GetInputDim("Scale").size(), 1UL);
  PADDLE_ENFORCE_EQ(ctx->GetInputDim("Scale")[0], C);
  PADDLE_ENFORCE_EQ(ctx->GetInputDim("Bias").size(), 1UL);
  PADDLE_ENFORCE_EQ(ctx->GetInputDim("Bias")[0], C);

  ctx->SetOutputDim("Y", x_dims);
  ctx->SetOutputDim("MeanOut", {C});
  ctx->SetOutputDim("VarianceOut", {C});
  ctx->SetOutputDim("SavedMean", {C});
  ctx->SetOutputDim("SavedVariance", {C});
  ctx->ShareLoD("X", "Y");
}

framework::OpKernelType BatchNormOp::GetExpectedKernelType(
    const framework::ExecutionContext &ctx) const {
  auto input_data_type = ctx.Input<Tensor>("X")->type();
  // By default, the type of the scale, bias, mean,
  // and var tensors should both be float. (For float or float16 input tensor)
  // or double (For double input tensor).
  auto bn_param_type = framework::proto::VarType::FP32;
  if (input_data_type == framework::proto::VarType::FP64) {
    bn_param_type = framework::proto::VarType::FP64;
  }
  PADDLE_ENFORCE_EQ(bn_param_type, ctx.Input<Tensor>("Scale")->type(),
                    "Scale input should be of float type");
  PADDLE_ENFORCE_EQ(bn_param_type, ctx.Input<Tensor>("Bias")->type(),
                    "Bias input should be of float type");
  PADDLE_ENFORCE_EQ(bn_param_type, ctx.Input<Tensor>("Mean")->type(),
                    "Mean input should be of float type");
  PADDLE_ENFORCE_EQ(bn_param_type, ctx.Input<Tensor>("Variance")->type(),
                    "Variance input should be of float type");

  // TODO(pzelazko-intel): enable MKLDNN layout when it's ready
  framework::LibraryType library = framework::LibraryType::kPlain;
  framework::DataLayout layout = framework::DataLayout::kAnyLayout;
103
#ifdef PADDLE_WITH_MKLDNN
Q
qingqing01 已提交
104 105 106 107
  if (library == framework::LibraryType::kPlain &&
      platform::CanMKLDNNBeUsed(ctx)) {
    library = framework::LibraryType::kMKLDNN;
    layout = framework::DataLayout::kMKLDNN;
K
Kexin Zhao 已提交
108
  }
Q
qingqing01 已提交
109
#endif
Q
Qiao Longfei 已提交
110

Q
qingqing01 已提交
111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170
  return framework::OpKernelType(input_data_type, ctx.GetPlace(), layout,
                                 library);
}

void BatchNormOpMaker::Make() {
  AddAttr<bool>("is_test",
                "(bool, default false) Set to true for inference only, false "
                "for training. Some layers may run faster when this is true.")
      .SetDefault(false);
  AddAttr<float>("momentum", "").SetDefault(0.9);
  AddAttr<float>("epsilon", "")
      .SetDefault(1e-5)
      .AddCustomChecker([](const float &epsilon) {
        PADDLE_ENFORCE(epsilon >= 0.0f && epsilon <= 0.001f,
                       "'epsilon' should be between 0.0 and 0.001.");
      });
  AddAttr<std::string>("data_layout", "").SetDefault("NCHW");
  AddInput("X", "The input tensor");
  AddInput("Scale",
           "Scale is a 1-dimensional tensor of size C "
           "that is applied to the output");
  AddInput("Bias",
           "Bias is a 1-dimensional tensor of size C "
           "that is applied to the output");
  AddInput("Mean",
           "The global mean (for training) or "
           "estimated mean (for testing)");
  AddInput("Variance",
           "The global variance (for training) "
           "or estimated Variance (for testing)");
  AddOutput("Y", "result after normalization");
  AddOutput("MeanOut",
            "Share memory with Mean. "
            "Store the global mean when training");
  AddOutput("VarianceOut",
            "Share memory with Variance. "
            "Store the global Variance when training");
  AddOutput("SavedMean",
            "Mean of the current mini batch, "
            "will apply to output when training")
      .AsIntermediate();
  AddOutput("SavedVariance",
            "Variance of the current mini batch, "
            "will apply to output when training")
      .AsIntermediate();
  AddAttr<bool>("use_mkldnn",
                "(bool, default false) Only used in mkldnn kernel")
      .SetDefault(false);
  AddAttr<bool>("fuse_with_relu",
                "(bool, default false) Only used in mkldnn kernel")
      .SetDefault(false);
  AddAttr<bool>("use_global_stats",
                "(bool, default false) Whether to use global mean and "
                "variance. In inference or test mode, set use_global_stats "
                "to true or is_test true. the behavior is equivalent. "
                "In train mode, when setting use_global_stats True, the "
                "global mean and variance are also used during train time, "
                "the BN acts as scaling and shiffting.")
      .SetDefault(false);
  AddComment(R"DOC(
171
Batch Normalization.
Q
Qiao Longfei 已提交
172

173 174 175 176 177 178
Batch Norm has been implemented as discussed in the paper:
https://arxiv.org/pdf/1502.03167.pdf
Can be used as a normalizer function for conv2d and fully_connected operations.
The required data format for this layer is one of the following:
1. NHWC `[batch, in_height, in_width, in_channels]`
2. NCHW `[batch, in_channels, in_height, in_width]`
Q
Qiao Longfei 已提交
179 180

)DOC");
Q
qingqing01 已提交
181
}
C
chengduo 已提交
182

Q
Qiao Longfei 已提交
183
template <typename T>
Q
QI JUN 已提交
184 185
class BatchNormKernel<platform::CPUDeviceContext, T>
    : public framework::OpKernel<T> {
Q
Qiao Longfei 已提交
186 187 188 189 190
 public:
  void Compute(const framework::ExecutionContext &ctx) const override {
    const float epsilon = ctx.Attr<float>("epsilon");
    const float momentum = ctx.Attr<float>("momentum");
    const bool is_test = ctx.Attr<bool>("is_test");
191 192 193 194
    const bool use_global_stats = ctx.Attr<bool>("use_global_stats");

    bool global_stats = is_test || use_global_stats;

Q
QI JUN 已提交
195 196 197
    const std::string data_layout_str = ctx.Attr<std::string>("data_layout");
    const DataLayout data_layout =
        framework::StringToDataLayout(data_layout_str);
Q
Qiao Longfei 已提交
198 199 200

    const auto *x = ctx.Input<Tensor>("X");
    const auto &x_dims = x->dims();
201 202
    PADDLE_ENFORCE(x_dims.size() >= 2 && x_dims.size() <= 5,
                   "The Input dim size should be between 2 and 5");
Q
Qiao Longfei 已提交
203 204
    const int N = x_dims[0];
    const int C =
Q
QI JUN 已提交
205 206
        (data_layout == DataLayout::kNCHW ? x_dims[1]
                                          : x_dims[x_dims.size() - 1]);
Q
Qiao Longfei 已提交
207 208 209 210 211 212 213 214 215 216 217 218 219 220 221
    const int sample_size = x->numel() / N / C;

    auto *y = ctx.Output<Tensor>("Y");
    auto *mean_out = ctx.Output<Tensor>("MeanOut");
    auto *variance_out = ctx.Output<Tensor>("VarianceOut");
    auto *saved_mean = ctx.Output<Tensor>("SavedMean");
    auto *saved_variance = ctx.Output<Tensor>("SavedVariance");

    // alloc memory
    y->mutable_data<T>(ctx.GetPlace());
    mean_out->mutable_data<T>(ctx.GetPlace());
    variance_out->mutable_data<T>(ctx.GetPlace());
    saved_mean->mutable_data<T>(ctx.GetPlace());
    saved_variance->mutable_data<T>(ctx.GetPlace());

222
    if (!global_stats) {
Q
Qiao Longfei 已提交
223 224 225 226 227 228 229 230
      // saved_xx is use just in this batch of data
      EigenVectorArrayMap<T> saved_mean_e(
          saved_mean->mutable_data<T>(ctx.GetPlace()), C);
      EigenVectorArrayMap<T> saved_variance_e(
          saved_variance->mutable_data<T>(ctx.GetPlace()), C);
      saved_mean_e.setZero();
      saved_variance_e.setZero();

231 232 233 234 235 236 237 238
      EigenVectorArrayMap<T> running_mean_arr(
          mean_out->mutable_data<T>(ctx.GetPlace()), C);
      EigenVectorArrayMap<T> running_var_arr(
          variance_out->mutable_data<T>(ctx.GetPlace()), C);

      if ((N * sample_size) == 1) {
        LOG(WARNING) << "Only 1 element in normalization dimension, "
                     << "we skip the batch norm calculation, let y = x.";
239
        framework::TensorCopy(*x, ctx.GetPlace(), y);
240 241 242
        return;
      }

Q
QI JUN 已提交
243 244
      switch (data_layout) {
        case DataLayout::kNCHW: {
Q
Qiao Longfei 已提交
245 246 247 248 249 250 251 252 253 254 255 256
          ConstEigenArrayMap<T> x_arr(x->data<T>(), sample_size, N * C);
          for (int nc = 0; nc < N * C; ++nc) {
            saved_mean_e(nc % C) += x_arr.col(nc).sum();
          }
          saved_mean_e /= N * sample_size;
          for (int nc = 0; nc < N * C; ++nc) {
            saved_variance_e(nc % C) +=
                (x_arr.col(nc) - saved_mean_e(nc % C)).matrix().squaredNorm();
          }
          saved_variance_e /= N * sample_size;
          break;
        }
Q
QI JUN 已提交
257
        case DataLayout::kNHWC: {
Q
Qiao Longfei 已提交
258 259 260 261 262 263 264 265 266 267 268 269 270
          ConstEigenArrayMap<T> x_arr(x->data<T>(), C, N * sample_size);
          for (int i = 0; i < N * sample_size; ++i) {
            saved_mean_e += x_arr.col(i);
          }
          saved_mean_e /= N * sample_size;
          for (int i = 0; i < N * sample_size; ++i) {
            saved_variance_e +=
                (x_arr.col(i) - saved_mean_e) * (x_arr.col(i) - saved_mean_e);
          }
          saved_variance_e /= N * sample_size;
          break;
        }
        default:
Q
QI JUN 已提交
271
          PADDLE_THROW("Unknown storage order: %s", data_layout_str);
Q
Qiao Longfei 已提交
272 273 274 275 276 277 278 279 280 281
      }

      running_mean_arr =
          running_mean_arr * momentum + saved_mean_e * (1. - momentum);
      running_var_arr =
          running_var_arr * momentum + saved_variance_e * (1. - momentum);
    }

    // use SavedMean and SavedVariance to do normalize
    Eigen::Array<T, Eigen::Dynamic, 1> inv_std(C);
282
    if (global_stats) {
Q
Qiao Longfei 已提交
283 284 285 286 287 288 289 290 291 292 293
      ConstEigenVectorArrayMap<T> var_arr(
          ctx.Input<Tensor>("Variance")->data<T>(), C);
      inv_std = (var_arr + epsilon).sqrt().inverse();
    } else {
      EigenVectorArrayMap<T> saved_inv_std(
          ctx.Output<Tensor>("SavedVariance")->data<T>(), C);
      // inverse SavedVariance first, gradient will use it too.
      saved_inv_std = (saved_inv_std + epsilon).inverse().sqrt();
      inv_std = saved_inv_std;
    }
    ConstEigenVectorArrayMap<T> mean_arr(
294 295
        global_stats ? ctx.Input<Tensor>("Mean")->data<T>()
                     : ctx.Output<Tensor>("SavedMean")->data<T>(),
Q
Qiao Longfei 已提交
296 297 298 299 300 301 302 303 304 305 306 307 308
        C);

    //   ((x - est_mean) * (inv_var) * scale + bias
    //   formula transform ====>
    //   (x * inv_var * scale) + (bias - est_mean * inv_var * scale)
    const auto *scale = ctx.Input<Tensor>("Scale");
    const auto *bias = ctx.Input<Tensor>("Bias");
    ConstEigenVectorArrayMap<T> scale_arr(scale->data<T>(), C);
    ConstEigenVectorArrayMap<T> bias_arr(bias->data<T>(), C);
    Eigen::Array<T, Eigen::Dynamic, 1> new_scale = inv_std * scale_arr;
    Eigen::Array<T, Eigen::Dynamic, 1> new_bias =
        bias_arr - mean_arr * inv_std * scale_arr;

Q
QI JUN 已提交
309 310
    switch (data_layout) {
      case DataLayout::kNCHW: {
Q
Qiao Longfei 已提交
311 312 313 314 315 316 317 318
        EigenArrayMap<T> y_arr(y->mutable_data<T>(ctx.GetPlace()), sample_size,
                               N * C);
        ConstEigenArrayMap<T> x_arr(x->data<T>(), sample_size, N * C);
        for (int nc = 0; nc < N * C; ++nc) {
          y_arr.col(nc) = x_arr.col(nc) * new_scale(nc % C) + new_bias(nc % C);
        }
        break;
      }
Q
QI JUN 已提交
319
      case DataLayout::kNHWC: {
Q
Qiao Longfei 已提交
320 321 322 323 324 325 326 327 328
        EigenArrayMap<T>(y->mutable_data<T>(ctx.GetPlace()), C,
                         N * sample_size) =
            (ConstEigenArrayMap<T>(x->data<T>(), C, N * sample_size).colwise() *
             new_scale)
                .colwise() +
            new_bias;
        break;
      }
      default:
Q
QI JUN 已提交
329
        PADDLE_THROW("Unknown storage order: %d", data_layout);
Q
Qiao Longfei 已提交
330 331 332 333
    }
  }
};

Q
qingqing01 已提交
334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357
void BatchNormGradOp::InferShape(framework::InferShapeContext *ctx) const {
  // check input
  PADDLE_ENFORCE(ctx->HasInput("X"));
  PADDLE_ENFORCE(ctx->HasInput("Scale"), "Input(scale) should not be null.");
  PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Y")),
                 "Input(Y@GRAD) should not be null.");
  PADDLE_ENFORCE(ctx->HasInput("SavedMean"),
                 "Input(SavedMean) should not be null.");
  PADDLE_ENFORCE(ctx->HasInput("SavedVariance"),
                 "Input(SavedVariance) should not be null");

  // check output
  PADDLE_ENFORCE(ctx->HasOutput(framework::GradVarName("X")), "");
  if (ctx->HasOutput(framework::GradVarName("Scale"))) {
    PADDLE_ENFORCE(ctx->HasOutput(framework::GradVarName("Bias")),
                   "Output(Scale@GRAD) and Output(Bias@GRAD) should not be "
                   "null at same time");
  }
  const bool use_global_stats = ctx->Attrs().Get<bool>("use_global_stats");
  if (use_global_stats) {
    PADDLE_ENFORCE(!ctx->Attrs().Get<bool>("use_mkldnn"),
                   "Using global stats during training is not supported "
                   "in gradient op kernel of batch_norm_mkldnn_op now.");
  }
Q
Qiao Longfei 已提交
358

Q
qingqing01 已提交
359 360 361 362 363
  const auto x_dims = ctx->GetInputDim("X");
  const DataLayout data_layout = framework::StringToDataLayout(
      ctx->Attrs().Get<std::string>("data_layout"));
  const int C = (data_layout == DataLayout::kNCHW ? x_dims[1]
                                                  : x_dims[x_dims.size() - 1]);
Q
Qiao Longfei 已提交
364

Q
qingqing01 已提交
365 366 367 368
  ctx->SetOutputDim(framework::GradVarName("X"), x_dims);
  if (ctx->HasOutput(framework::GradVarName("Scale"))) {
    ctx->SetOutputDim(framework::GradVarName("Scale"), {C});
    ctx->SetOutputDim(framework::GradVarName("Bias"), {C});
Q
Qiao Longfei 已提交
369
  }
Q
qingqing01 已提交
370
}
Q
Qiao Longfei 已提交
371

Q
qingqing01 已提交
372 373 374 375 376 377 378 379 380 381 382 383 384 385 386
framework::OpKernelType BatchNormGradOp::GetExpectedKernelType(
    const framework::ExecutionContext &ctx) const {
  const auto *var = ctx.InputVar(framework::GradVarName("Y"));
  if (var == nullptr) {
    PADDLE_THROW("can't find Y@GRAD");
  }
  const Tensor *t = nullptr;
  if (var->IsType<Tensor>()) {
    t = &var->Get<Tensor>();
  } else if (var->IsType<LoDTensor>()) {
    t = &var->Get<LoDTensor>();
  }
  if (t == nullptr) {
    PADDLE_THROW("can't find Y@GRAD");
  }
387

Q
qingqing01 已提交
388 389 390
  // TODO(pzelazko-intel): enable MKLDNN layout when it's ready
  framework::LibraryType library = framework::LibraryType::kPlain;
  framework::DataLayout layout = framework::DataLayout::kAnyLayout;
391

392
#ifdef PADDLE_WITH_MKLDNN
Q
qingqing01 已提交
393 394 395 396 397
  if (library == framework::LibraryType::kPlain &&
      platform::CanMKLDNNBeUsed(ctx)) {
    library = framework::LibraryType::kMKLDNN;
    layout = framework::DataLayout::kMKLDNN;
  }
398
#endif
399

Q
qingqing01 已提交
400 401 402
  return framework::OpKernelType(ctx.Input<Tensor>("X")->type(), ctx.GetPlace(),
                                 layout, library);
}
Q
Qiao Longfei 已提交
403 404

template <typename T>
Q
QI JUN 已提交
405
class BatchNormGradKernel<platform::CPUDeviceContext, T>
Q
Qiao Longfei 已提交
406 407 408 409 410 411 412 413 414
    : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext &ctx) const override {
    const auto *x = ctx.Input<Tensor>("X");
    const auto *d_y = ctx.Input<Tensor>(framework::GradVarName("Y"));
    const auto *scale = ctx.Input<Tensor>("Scale");
    const auto *saved_mean = ctx.Input<Tensor>("SavedMean");
    // SavedVariance have been reverted in forward operator
    const auto *saved_inv_variance = ctx.Input<Tensor>("SavedVariance");
Q
QI JUN 已提交
415
    const std::string data_layout_str = ctx.Attr<std::string>("data_layout");
416 417
    const bool use_global_stats = ctx.Attr<bool>("use_global_stats");
    const float epsilon = ctx.Attr<float>("epsilon");
Q
QI JUN 已提交
418 419
    const DataLayout data_layout =
        framework::StringToDataLayout(data_layout_str);
Q
Qiao Longfei 已提交
420 421 422 423

    // Get the size for each dimension.
    // NCHW [batch_size, in_channels, in_height, in_width]
    const auto &x_dims = x->dims();
424 425
    PADDLE_ENFORCE(x_dims.size() >= 2 && x_dims.size() <= 5,
                   "The Input dim size should be between 2 and 5");
Q
Qiao Longfei 已提交
426 427
    const int N = x_dims[0];
    const int C =
Q
QI JUN 已提交
428 429
        (data_layout == DataLayout::kNCHW ? x_dims[1]
                                          : x_dims[x_dims.size() - 1]);
Q
Qiao Longfei 已提交
430 431 432 433 434 435 436 437
    const int sample_size = x->numel() / N / C;

    // init output
    auto *d_x = ctx.Output<Tensor>(framework::GradVarName("X"));
    auto *d_scale = ctx.Output<Tensor>(framework::GradVarName("Scale"));
    auto *d_bias = ctx.Output<Tensor>(framework::GradVarName("Bias"));

    d_x->mutable_data<T>(ctx.GetPlace());
438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465

    const T *mean_data = saved_mean->data<T>();
    const T *inv_var_data = saved_inv_variance->data<T>();
    Tensor inv_var_tensor;
    if (use_global_stats) {
      const auto *running_mean = ctx.Input<Tensor>("Mean");
      const auto *running_variance = ctx.Input<Tensor>("Variance");
      mean_data = running_mean->data<T>();
      T *running_inv_var_data = inv_var_tensor.mutable_data<T>(ctx.GetPlace());
      EigenVectorArrayMap<T> inv_var_tmp(running_inv_var_data, C);
      ConstEigenVectorArrayMap<T> var_arr(running_variance->data<T>(), C);

      inv_var_tmp = (var_arr + epsilon).sqrt().inverse().eval();
      inv_var_data = running_inv_var_data;
    }

    ConstEigenVectorArrayMap<T> scale_arr(scale->data<T>(), C);
    ConstEigenVectorArrayMap<T> mean_arr(mean_data, C);
    ConstEigenVectorArrayMap<T> inv_var_arr(inv_var_data, C);

    T *d_bias_data = nullptr;
    T *d_scale_data = nullptr;
    if (d_scale && d_bias) {
      d_scale->mutable_data<T>(ctx.GetPlace());
      d_bias->mutable_data<T>(ctx.GetPlace());
      d_bias_data = d_bias->mutable_data<T>(ctx.GetPlace());
      d_scale_data = d_scale->mutable_data<T>(ctx.GetPlace());
    }
Q
Qiao Longfei 已提交
466 467 468 469 470

    // d_bias = np.sum(d_y, axis=0)
    // d_scale = np.sum((X - mean) / inv_std * dy, axis=0)
    // d_x = (1. / N) * scale * inv_var * (N * d_y - np.sum(d_y, axis=0)
    //   - (X - mean) * inv_var * inv_var * np.sum(d_y * (X - mean), axis=0))
471 472
    EigenVectorArrayMap<T> d_bias_arr(d_bias_data, C);
    EigenVectorArrayMap<T> d_scale_arr(d_scale_data, C);
Q
Qiao Longfei 已提交
473

474 475 476 477
    if (d_scale && d_bias) {
      d_bias_arr.setZero();
      d_scale_arr.setZero();
    }
Q
Qiao Longfei 已提交
478

479 480
    if ((N * sample_size) == 1 && !use_global_stats) {
      framework::TensorCopy(*d_y, ctx.GetPlace(), d_x);
481 482 483
      return;
    }

484 485
    int scale_coefff = use_global_stats ? 1 : N * sample_size;
    const auto scale_inv_var_nhw = scale_arr * inv_var_arr / scale_coefff;
Q
Qiao Longfei 已提交
486

Q
QI JUN 已提交
487 488
    switch (data_layout) {
      case DataLayout::kNCHW: {
Q
Qiao Longfei 已提交
489 490 491 492 493 494
        ConstEigenArrayMap<T> x_arr(x->data<T>(), sample_size, N * C);
        ConstEigenArrayMap<T> d_y_arr(d_y->data<T>(), sample_size, N * C);
        EigenArrayMap<T> d_x_arr(d_x->mutable_data<T>(ctx.GetPlace()),
                                 sample_size, N * C);
        d_x_arr.setZero();

495 496 497 498 499 500 501 502
        if (d_scale && d_bias) {
          for (int nc = 0; nc < N * C; ++nc) {
            int c = nc % C;
            d_bias_arr(c) += d_y_arr.col(nc).sum();
            d_scale_arr(c) += ((x_arr.col(nc) - mean_arr(c)) * inv_var_arr(c) *
                               d_y_arr.col(nc))
                                  .sum();
          }
Q
Qiao Longfei 已提交
503
        }
504 505 506 507 508 509 510 511 512 513 514 515 516 517
        if (!use_global_stats) {
          for (int nc = 0; nc < N * C; ++nc) {
            int c = nc % C;
            d_x_arr.col(nc) +=
                scale_inv_var_nhw(c) *
                (d_y_arr.col(nc) * N * sample_size - d_bias_arr(c) -
                 (x_arr.col(nc) - mean_arr[c]) * d_scale_arr(c) *
                     inv_var_arr(c));
          }
        } else {
          for (int nc = 0; nc < N * C; ++nc) {
            int c = nc % C;
            d_x_arr.col(nc) += scale_inv_var_nhw(c) * d_y_arr.col(nc);
          }
Q
Qiao Longfei 已提交
518 519 520
        }
        break;
      }
Q
QI JUN 已提交
521
      case DataLayout::kNHWC: {
Q
Qiao Longfei 已提交
522 523 524 525 526 527 528 529 530 531 532
        ConstEigenArrayMap<T> x_arr(x->data<T>(), C, N * sample_size);
        ConstEigenArrayMap<T> d_y_arr(d_y->data<T>(), C, N * sample_size);
        EigenArrayMap<T> d_x_arr(d_x->mutable_data<T>(ctx.GetPlace()), C,
                                 N * sample_size);
        d_x_arr.setZero();

        const auto d_y_row_sum = d_y_arr.rowwise().sum();
        const auto x_minus_mean = x_arr.colwise() - mean_arr;
        const auto d_y_mul_x_minus_mean_row_sum =
            (d_y_arr * x_minus_mean).rowwise().sum();
        const auto inv_var_sqr = inv_var_arr * inv_var_arr;
533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553

        if (d_scale && d_bias) {
          for (int nhw = 0; nhw < N * sample_size; ++nhw) {
            d_bias_arr += d_y_arr.col(nhw);
            d_scale_arr +=
                (x_arr.col(nhw) - mean_arr) * inv_var_arr * d_y_arr.col(nhw);
          }
        }

        if (!use_global_stats) {
          for (int nhw = 0; nhw < N * sample_size; ++nhw) {
            d_x_arr.col(nhw) +=
                scale_inv_var_nhw *
                (d_y_arr.col(nhw) * N * sample_size - d_y_row_sum -
                 x_minus_mean.col(nhw) * inv_var_sqr *
                     d_y_mul_x_minus_mean_row_sum);
          }
        } else {
          for (int nhw = 0; nhw < N * sample_size; ++nhw) {
            d_x_arr.col(nhw) += scale_inv_var_nhw * d_y_arr.col(nhw);
          }
Q
Qiao Longfei 已提交
554 555 556 557
        }
        break;
      }
      default:
Q
QI JUN 已提交
558
        PADDLE_THROW("Unknown storage order: %s", data_layout_str);
Q
Qiao Longfei 已提交
559 560 561 562
    }
  }
};

Q
qingqing01 已提交
563 564 565 566 567 568 569
std::unique_ptr<framework::OpDesc> BatchNormGradMaker::Apply() const {
  auto *op = new framework::OpDesc();
  op->SetType(GradOpType());
  op->SetInput("X", Input("X"));
  op->SetInput(framework::GradVarName("Y"), OutputGrad("Y"));

  op->SetInput("Scale", Input("Scale"));
S
sneaxiy 已提交
570
  op->SetInput("Bias", Input("Bias"));
Q
qingqing01 已提交
571 572 573 574 575 576 577 578
  op->SetInput("SavedMean", Output("SavedMean"));
  op->SetInput("SavedVariance", Output("SavedVariance"));

  // used when setting use_global_stats True during training
  if (boost::get<bool>(GetAttr("use_global_stats"))) {
    op->SetInput("Mean", Output("MeanOut"));
    op->SetInput("Variance", Output("VarianceOut"));
  }
579

Q
qingqing01 已提交
580
  op->SetAttrMap(Attrs());
Y
Yu Yang 已提交
581

Q
qingqing01 已提交
582 583 584
  op->SetOutput(framework::GradVarName("X"), InputGrad("X"));
  op->SetOutput(framework::GradVarName("Scale"), InputGrad("Scale"));
  op->SetOutput(framework::GradVarName("Bias"), InputGrad("Bias"));
Y
Yu Yang 已提交
585

Q
qingqing01 已提交
586 587
  return std::unique_ptr<framework::OpDesc>(op);
}
Y
Yu Yang 已提交
588

D
dzhwinter 已提交
589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621
class BatchNormInplaceInToOut : public framework::InplaceInToOut {
 public:
  using InplaceInToOut::InplaceInToOut;

 protected:
  std::unordered_map<std::string, std::string> Apply(
      const framework::OpDesc &op_desc,
      framework::BlockDesc *block) const override {
    std::unordered_map<std::string, std::string> inplace_in_to_out = {
        {"Mean", "MeanOut"}, {"Variance", "VarianceOut"}, {"X", "Y"},
    };
    return inplace_in_to_out;
  }
};

class BatchNormGradInplaceInToOut : public framework::InplaceInToOut {
 public:
  using InplaceInToOut::InplaceInToOut;

 protected:
  std::unordered_map<std::string, std::string> Apply(
      const framework::OpDesc &op_desc,
      framework::BlockDesc *block) const override {
    std::unordered_map<std::string, std::string> inplace_in_to_out = {
        // Scale, Bias, SavedMean, SavedVariance shape is [batch_size, C]
        {framework::GradVarName("Y"), framework::GradVarName("X")},
        {"SavedMean", framework::GradVarName("Scale")},
        {"SavedVariance", framework::GradVarName("Bias")},
    };
    return inplace_in_to_out;
  }
};

Q
Qiao Longfei 已提交
622 623 624 625
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
Y
Yu Yang 已提交
626
REGISTER_OPERATOR(batch_norm, ops::BatchNormOp, ops::BatchNormOpMaker,
Q
qingqing01 已提交
627 628 629 630
                  ops::BatchNormOpInferVarType, ops::BatchNormGradMaker)
// ops::BatchNormInplaceInToOut);
REGISTER_OPERATOR(batch_norm_grad, ops::BatchNormGradOp)
//                  ops::BatchNormGradInplaceInToOut);
Y
Yu Yang 已提交
631

Q
QI JUN 已提交
632
REGISTER_OP_CPU_KERNEL(
D
dzhwinter 已提交
633 634
    batch_norm, ops::BatchNormKernel<paddle::platform::CPUDeviceContext, float>,
    ops::BatchNormKernel<paddle::platform::CPUDeviceContext, double>);
Q
Qiao Longfei 已提交
635 636
REGISTER_OP_CPU_KERNEL(
    batch_norm_grad,
D
dzhwinter 已提交
637 638
    ops::BatchNormGradKernel<paddle::platform::CPUDeviceContext, float>,
    ops::BatchNormGradKernel<paddle::platform::CPUDeviceContext, double>);