#ifndef CAFFE2_OPERATORS_SPATIAL_BATCH_NORM_OP_H_
|
#define CAFFE2_OPERATORS_SPATIAL_BATCH_NORM_OP_H_
|
|
#include <algorithm>
|
#include <array>
|
#include <functional>
|
#include <string>
|
#include <vector>
|
|
#include "caffe2/core/context.h"
|
#include "caffe2/core/operator.h"
|
#include "caffe2/utils/eigen_utils.h"
|
#include "caffe2/utils/math.h"
|
|
namespace caffe2 {
|
|
template <class Context>
|
class SpatialBNOp : public Operator<Context> {
|
public:
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
|
template <class... Args>
|
explicit SpatialBNOp(Args&&... args)
|
: Operator<Context>(std::forward<Args>(args)...),
|
OP_SINGLE_ARG(bool, OpSchema::Arg_IsTest, is_test_, false),
|
OP_SINGLE_ARG(double, "epsilon", epsilon_, 1e-5),
|
OP_SINGLE_ARG(float, "momentum", momentum_, 0.9f),
|
order_(StringToStorageOrder(
|
this->template GetSingleArgument<std::string>("order", "NCHW"))),
|
OP_SINGLE_ARG(int, "num_batches", num_batches_, 1) {
|
CAFFE_ENFORCE_NE(
|
order_,
|
StorageOrder::UNKNOWN,
|
"order should be either \"NCHW\" or \"NHWC\".");
|
CAFFE_ENFORCE(
|
(is_test_ && OutputSize() == 1) || (!is_test_ && OutputSize() == 5));
|
CAFFE_ENFORCE_GT(epsilon_, 0);
|
CAFFE_ENFORCE_GE(momentum_, 0);
|
CAFFE_ENFORCE_LE(momentum_, 1);
|
}
|
|
virtual ~SpatialBNOp() = default;
|
|
bool RunOnDevice() override {
|
return DispatchHelper<TensorTypes<float>>::call(this, Input(0));
|
}
|
|
template <typename T>
|
bool DoRunWithType() {
|
const auto& X = Input(INPUT);
|
const auto& scale = Input(SCALE);
|
const auto& bias = Input(BIAS);
|
|
const int ndim = X.dim();
|
CAFFE_ENFORCE_GE(ndim, 3);
|
const int N = X.dim32(0);
|
const int C =
|
(order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(ndim - 1));
|
const std::vector<int> X_dims(X.sizes().cbegin(), X.sizes().cend());
|
const int HxW =
|
std::accumulate(
|
X_dims.cbegin() + 1, X_dims.cend(), 1, std::multiplies<int>()) /
|
C;
|
CAFFE_ENFORCE_EQ(scale.numel(), C);
|
CAFFE_ENFORCE_EQ(bias.numel(), C);
|
|
auto* Y = Output(OUTPUT, X.sizes(), at::dtype<T>());
|
const T* X_data = X.template data<T>();
|
const T* scale_data = scale.template data<T>();
|
const T* bias_data = bias.template data<T>();
|
T* Y_data = Y->template mutable_data<T>();
|
ReinitializeTensor(
|
&alpha_, {C}, at::dtype<T>().device(Context::GetDeviceType()));
|
ReinitializeTensor(
|
&beta_, {C}, at::dtype<T>().device(Context::GetDeviceType()));
|
T* alpha_data = alpha_.template mutable_data<T>();
|
T* beta_data = beta_.template mutable_data<T>();
|
if (is_test_) {
|
if (N == 0) {
|
return true;
|
}
|
const auto& mean = Input(EST_MEAN);
|
const auto& var = Input(EST_VAR);
|
CAFFE_ENFORCE_EQ(mean.numel(), C);
|
CAFFE_ENFORCE_EQ(var.numel(), C);
|
ComputeFusedParam<T>(
|
C,
|
scale_data,
|
bias_data,
|
mean.template data<T>(),
|
var.template data<T>(),
|
alpha_data,
|
beta_data);
|
} else {
|
auto* saved_mean = Output(SAVED_MEAN, {C}, at::dtype<T>());
|
auto* saved_rstd = Output(SAVED_INV_STD, {C}, at::dtype<T>());
|
T* saved_mean_data = saved_mean->template mutable_data<T>();
|
T* saved_rstd_data = saved_rstd->template mutable_data<T>();
|
|
// Enforce Alias
|
CAFFE_ENFORCE(
|
IsInputOutputAlias(3, 1), "Input 3 and Output 1 should be alias.");
|
CAFFE_ENFORCE(
|
IsInputOutputAlias(4, 2), "Input 4 and Output 2 should be alias.");
|
|
Tensor* running_mean = nullptr;
|
Tensor* running_var = nullptr;
|
const auto& mean = Input(EST_MEAN);
|
const auto& var = Input(EST_VAR);
|
if (mean.numel() != C) {
|
running_mean = Output(RUNNING_MEAN, {C}, at::dtype<T>());
|
C10_LOG_EVERY_MS(WARNING, 1000)
|
<< "[Depreacated] Running mean is not initialized in "
|
"SpatialBatchNorm Op";
|
math::Set<T, Context>(
|
C, T(0), running_mean->template mutable_data<T>(), &context_);
|
} else {
|
running_mean = Output(RUNNING_MEAN, {C}, at::dtype<T>());
|
}
|
if (var.numel() != C) {
|
running_var = Output(RUNNING_VAR, {C}, at::dtype<T>());
|
math::Set<T, Context>(
|
C, T(0), running_var->template mutable_data<T>(), &context_);
|
C10_LOG_EVERY_MS(WARNING, 1000)
|
<< "[Deprecated] Running variance is not initialized in "
|
"SpatialBatchNorm Op";
|
} else {
|
running_var = Output(RUNNING_VAR, {C}, at::dtype<T>());
|
}
|
|
T* running_mean_data = running_mean->template mutable_data<T>();
|
T* running_var_data = running_var->template mutable_data<T>();
|
if (N == 0) {
|
math::Set<T, Context>(C, T(0), saved_mean_data, &context_);
|
math::Set<T, Context>(C, T(0), saved_rstd_data, &context_);
|
return true;
|
}
|
if (num_batches_ > 1) {
|
const auto& batch_mean_sum = Input(BATCH_MEAN_SUM);
|
const auto& batch_var_sum = Input(BATCH_VAR_SUM);
|
CAFFE_ENFORCE_EQ(batch_mean_sum.numel(), C);
|
CAFFE_ENFORCE_EQ(batch_var_sum.numel(), C);
|
ComputeBatchMoments<T>(
|
N,
|
C,
|
HxW,
|
batch_mean_sum.template data<T>(),
|
batch_var_sum.template data<T>(),
|
saved_mean_data,
|
saved_rstd_data);
|
} else {
|
if (order_ == StorageOrder::NCHW) {
|
const std::array<int, 3> X_dims_arr = {N, C, HxW};
|
const std::array<int, 3> Y_dims_arr = {1, C, 1};
|
math::Moments<T, Context>(
|
3,
|
X_dims_arr.data(),
|
Y_dims_arr.data(),
|
X_data,
|
saved_mean_data,
|
saved_rstd_data,
|
&context_);
|
} else {
|
const std::array<int, 2> X_dims_arr = {N * HxW, C};
|
const std::array<int, 2> Y_dims_arr = {1, C};
|
math::Moments<T, Context>(
|
2,
|
X_dims_arr.data(),
|
Y_dims_arr.data(),
|
X_data,
|
saved_mean_data,
|
saved_rstd_data,
|
&context_);
|
}
|
}
|
ComputeRunningMomentsAndFusedParam<T>(
|
C,
|
scale_data,
|
bias_data,
|
saved_mean_data,
|
saved_rstd_data,
|
running_mean_data,
|
running_var_data,
|
saved_rstd_data,
|
alpha_data,
|
beta_data);
|
}
|
if (order_ == StorageOrder::NCHW) {
|
math::AffineChannel<T, Context, StorageOrder::NCHW>(
|
N, C, HxW, X_data, alpha_data, beta_data, Y_data, &context_);
|
} else {
|
math::AffineChannel<T, Context, StorageOrder::NHWC>(
|
N, C, HxW, X_data, alpha_data, beta_data, Y_data, &context_);
|
}
|
|
return true;
|
}
|
|
protected:
|
template <typename T>
|
void ComputeFusedParam(
|
const int C,
|
const T* scale,
|
const T* bias,
|
const T* mean,
|
const T* var,
|
T* alpha,
|
T* beta) {
|
EigenVectorArrayMap<T> alpha_arr(alpha, C);
|
EigenVectorArrayMap<T> beta_arr(beta, C);
|
alpha_arr = ConstEigenVectorArrayMap<T>(scale, C) *
|
(ConstEigenVectorArrayMap<T>(var, C) + static_cast<T>(epsilon_))
|
.rsqrt();
|
beta_arr = ConstEigenVectorArrayMap<T>(bias, C) -
|
alpha_arr * ConstEigenVectorArrayMap<T>(mean, C);
|
}
|
|
template <typename T>
|
void ComputeBatchMoments(
|
const int N,
|
const int C,
|
const int HxW,
|
const T* batch_mean_sum,
|
const T* batch_var_sum,
|
T* mean,
|
T* var) {
|
const T scale = T(1) / static_cast<T>(num_batches_ * N * HxW);
|
EigenVectorArrayMap<T> mean_arr(mean, C);
|
EigenVectorArrayMap<T> var_arr(var, C);
|
mean_arr = ConstEigenVectorArrayMap<T>(batch_mean_sum, C) * scale;
|
var_arr = ConstEigenVectorArrayMap<T>(batch_var_sum, C) * scale -
|
mean_arr.square();
|
}
|
|
template <typename T>
|
void ComputeRunningMomentsAndFusedParam(
|
const int C,
|
const T* scale,
|
const T* bias,
|
const T* mean,
|
const T* var,
|
T* running_mean,
|
T* running_var,
|
T* rstd,
|
T* alpha,
|
T* beta) {
|
const T a = T(1) - static_cast<T>(momentum_);
|
const T b = static_cast<T>(momentum_);
|
math::Axpby<T, T, Context>(C, a, mean, b, running_mean, &context_);
|
math::Axpby<T, T, Context>(C, a, var, b, running_var, &context_);
|
math::InvStd<T, Context>(C, static_cast<T>(epsilon_), var, rstd, &context_);
|
EigenVectorArrayMap<T> alpha_arr(alpha, C);
|
EigenVectorArrayMap<T> beta_arr(beta, C);
|
alpha_arr = ConstEigenVectorArrayMap<T>(scale, C) *
|
ConstEigenVectorArrayMap<T>(rstd, C);
|
beta_arr = ConstEigenVectorArrayMap<T>(bias, C) -
|
alpha_arr * ConstEigenVectorArrayMap<T>(mean, C);
|
}
|
|
const bool is_test_;
|
double epsilon_;
|
const float momentum_;
|
const StorageOrder order_;
|
const int num_batches_;
|
|
Tensor alpha_;
|
Tensor beta_;
|
|
INPUT_TAGS(
|
INPUT,
|
SCALE,
|
BIAS,
|
EST_MEAN,
|
EST_VAR,
|
BATCH_MEAN_SUM,
|
BATCH_VAR_SUM);
|
OUTPUT_TAGS(OUTPUT, RUNNING_MEAN, RUNNING_VAR, SAVED_MEAN, SAVED_INV_STD);
|
};
|
|
template <class Context>
|
class SpatialBNGradientOp : public Operator<Context> {
|
public:
|
USE_OPERATOR_CONTEXT_FUNCTIONS;
|
|
template <class... Args>
|
explicit SpatialBNGradientOp(Args&&... args)
|
: Operator<Context>(std::forward<Args>(args)...),
|
OP_SINGLE_ARG(double, "epsilon", epsilon_, 1e-5),
|
order_(StringToStorageOrder(
|
this->template GetSingleArgument<string>("order", "NCHW"))),
|
OP_SINGLE_ARG(int, "num_batches", num_batches_, 1) {
|
CAFFE_ENFORCE_NE(
|
order_,
|
StorageOrder::UNKNOWN,
|
"order should be either \"NCHW\" or \"NHWC\".");
|
CAFFE_ENFORCE(InputSize() == 5 || InputSize() == 7);
|
CAFFE_ENFORCE_EQ(OutputSize(), 3);
|
}
|
|
virtual ~SpatialBNGradientOp() = default;
|
|
bool RunOnDevice() override {
|
return DispatchHelper<TensorTypes<float>>::call(this, Input(0));
|
}
|
|
template <typename T>
|
bool DoRunWithType() {
|
const auto& X = Input(INPUT);
|
const auto& dY = Input(OUTPUT_GRAD);
|
const auto& scale = Input(SCALE);
|
const auto& mean = Input(SAVED_MEAN);
|
const auto& rstd = Input(SAVED_INV_STD);
|
const int ndim = X.dim();
|
CAFFE_ENFORCE_GE(ndim, 3);
|
const int N = X.dim32(0);
|
const int C =
|
(order_ == StorageOrder::NCHW ? X.dim32(1) : X.dim32(ndim - 1));
|
const std::vector<int> X_dims(X.sizes().cbegin(), X.sizes().cend());
|
const int HxW =
|
std::accumulate(
|
X_dims.cbegin() + 1, X_dims.cend(), 1, std::multiplies<int>()) /
|
C;
|
CAFFE_ENFORCE_EQ(scale.numel(), C);
|
CAFFE_ENFORCE_EQ(mean.numel(), C);
|
CAFFE_ENFORCE_EQ(rstd.numel(), C);
|
|
auto* dX = Output(INPUT_GRAD, X.sizes(), at::dtype<T>());
|
at::IntArrayRef dscale_sizes, dbias_sizes;
|
if (num_batches_ == 1) {
|
dscale_sizes = scale.sizes();
|
dbias_sizes = scale.sizes();
|
} else {
|
const auto& dscale_sum = Input(AGGREGATE_SCALE_GRAD);
|
const auto& dbias_sum = Input(AGGREGATE_BIAS_GRAD);
|
// Note: previously there was alias check to decide whether to call
|
// ResizeLike or not, since we only call Resize when the size does not
|
// match the size of cached Tensor, this check is not necessary
|
dscale_sizes = dscale_sum.sizes();
|
dbias_sizes = dbias_sum.sizes();
|
}
|
auto* dscale = Output(SCALE_GRAD, dscale_sizes, at::dtype<T>());
|
auto* dbias = Output(BIAS_GRAD, dbias_sizes, at::dtype<T>());
|
const T* X_data = X.template data<T>();
|
const T* dY_data = dY.template data<T>();
|
const T* scale_data = scale.template data<T>();
|
const T* mean_data = mean.template data<T>();
|
const T* rstd_data = rstd.template data<T>();
|
T* dX_data = dX->template mutable_data<T>();
|
T* dscale_data = dscale->template mutable_data<T>();
|
T* dbias_data = dbias->template mutable_data<T>();
|
|
if (N == 0) {
|
math::Set<T, Context>(C, T(0), dscale_data, &context_);
|
math::Set<T, Context>(C, T(0), dbias_data, &context_);
|
return true;
|
}
|
ReinitializeTensor(
|
&alpha_, {C}, at::dtype<T>().device(Context::GetDeviceType()));
|
ReinitializeTensor(
|
&beta_, {C}, at::dtype<T>().device(Context::GetDeviceType()));
|
ReinitializeTensor(
|
&gamma_, {C}, at::dtype<T>().device(Context::GetDeviceType()));
|
T* alpha_data = alpha_.template mutable_data<T>();
|
T* beta_data = beta_.template mutable_data<T>();
|
T* gamma_data = gamma_.template mutable_data<T>();
|
if (num_batches_ > 1) {
|
const auto& dscale_sum = Input(AGGREGATE_SCALE_GRAD);
|
const auto& dbias_sum = Input(AGGREGATE_BIAS_GRAD);
|
ComputeMultiBatchScaleBiasGradientsAndFusedParams<T>(
|
N,
|
C,
|
HxW,
|
scale_data,
|
mean_data,
|
rstd_data,
|
dscale_sum.template data<T>(),
|
dbias_sum.template data<T>(),
|
dscale_data,
|
dbias_data,
|
alpha_data,
|
beta_data,
|
gamma_data);
|
} else {
|
ComputeScaleBiasGradientsAndFusedParams<T>(
|
N,
|
C,
|
HxW,
|
dY_data,
|
X_data,
|
scale_data,
|
mean_data,
|
rstd_data,
|
dscale_data,
|
dbias_data,
|
alpha_data,
|
beta_data,
|
gamma_data,
|
dX_data);
|
}
|
ComputeXGradient<T>(
|
N, C, HxW, dY_data, X_data, alpha_data, beta_data, gamma_data, dX_data);
|
|
return true;
|
}
|
|
protected:
|
template <typename T>
|
void ComputeMultiBatchScaleBiasGradientsAndFusedParams(
|
const int N,
|
const int C,
|
const int HxW,
|
const T* scale,
|
const T* mean,
|
const T* rstd,
|
const T* dscale_sum,
|
const T* dbias_sum,
|
T* dscale,
|
T* dbias,
|
T* alpha,
|
T* beta,
|
T* gamma);
|
|
template <typename T>
|
void ComputeScaleBiasGradientsAndFusedParams(
|
const int N,
|
const int C,
|
const int HxW,
|
const T* dY,
|
const T* X,
|
const T* scale,
|
const T* mean,
|
const T* rstd,
|
T* dscale,
|
T* dbias,
|
T* alpha,
|
T* beta,
|
T* gamma,
|
T* scratch);
|
|
template <typename T>
|
void ComputeXGradient(
|
const int N,
|
const int C,
|
const int HxW,
|
const T* dY,
|
const T* X,
|
const T* alpha,
|
const T* beta,
|
const T* gamma,
|
T* dX);
|
|
double epsilon_;
|
const StorageOrder order_;
|
const int num_batches_;
|
|
Tensor alpha_;
|
Tensor beta_;
|
Tensor gamma_;
|
Tensor ones_;
|
|
INPUT_TAGS(
|
INPUT,
|
SCALE,
|
OUTPUT_GRAD,
|
SAVED_MEAN,
|
SAVED_INV_STD,
|
AGGREGATE_SCALE_GRAD,
|
AGGREGATE_BIAS_GRAD);
|
OUTPUT_TAGS(INPUT_GRAD, SCALE_GRAD, BIAS_GRAD);
|
};
|
|
} // namespace caffe2
|
|
#endif // CAFFE2_OPERATORS_SPATIAL_BATCH_NORM_OP_H_
|
