Method BatchNormHandle

src/model/operation/batchnorm.cc:27–64

namespace singa {

BatchNormHandle::BatchNormHandle(const float momentum, const Tensor& input) {
  factor = momentum;
  batchsize = input.shape(0);
  channels = input.shape(1);
  if (input.nDim() == 4u) {
    height = input.shape().at(2);
    width = input.shape().at(3);
    is_2d = false;
  } else if (input.nDim() == 2u) {
    height = 1;
    width = 1;
    is_2d = true;
  } else {
    LOG(FATAL) << "The dimension of input should either be 4D or 2D.";
  }

#ifdef USE_DNNL
  if (input.device()->lang() == kCpp) {
    use_dnnl = true;
    epsilon = 1e-5f;
    x_dims = dnnl::memory::dims(input.shape().begin(), input.shape().end());

    // support f32 only
    auto dtype_ = memory::data_type::f32;
    memory::format_tag format_tag_ = get_dnnl_format_tag(input);
    x_md = dnnl::memory::desc({x_dims}, dtype_, format_tag_);

    // add to
    bn_fwd_training_d = new dnnl::batch_normalization_forward::desc(
        dnnl::prop_kind::forward_training, x_md, epsilon,
        dnnl::normalization_flags::use_scale_shift);

    auto eng = input.device()->context(0)->dnnl_engine;
    bn_fwd_training_pd = new dnnl::batch_normalization_forward::primitive_desc(
        *bn_fwd_training_d, eng);
  }
#endif  // USE_DNNL
};

BatchNormHandle::~BatchNormHandle() {
#ifdef USE_DNNL