spatial_batch_norm_dnnlowp_op.h

#pragma once

#include "caffe2/operators/spatial_batch_norm_op.h"
#include "caffe2/quantization/server/dnnlowp_op.h"

namespace caffe2 {

template <typename T>
class SpatialBNDNNLowPOp final : public DNNLowPOp<T, SpatialBNOp<CPUContext>> {
 public:
  USE_OPERATOR_FUNCTIONS(CPUContext);
  USE_DNNLOWP_OPERATOR_BASE_FUNCTIONS(T, SpatialBNOp<CPUContext>);
  SpatialBNDNNLowPOp(const OperatorDef& operator_def, Workspace* ws);

  virtual ~SpatialBNDNNLowPOp() override = default;

  bool RunOnDevice() override;

 private:
  void ComputeFusedParam_(
      const int C,
      const float* scale,
      const float* bias,
      const float* mean,
      const float* var,
      float* alpha,
      float* beta);

  double epsilon_;
  const StorageOrder order_;

  Tensor alpha_;
  Tensor beta_;

  INPUT_TAGS(INPUT, SCALE, BIAS, EST_MEAN, EST_VAR);
  OUTPUT_TAGS(OUTPUT);
};

} // namespace caffe2