in src/inplace_abn_cpu.cpp [107:142]
void forward_cpu(
at::Tensor& x_,
const at::Tensor& mean,
const at::Tensor& var,
const c10::optional<at::Tensor>& weight,
const c10::optional<at::Tensor>& bias,
float eps,
Activation activation,
float activation_param) {
CHECK_NOT_HALF(x_);
auto x = normalize_shape(x_);
// Apply normalization
auto abs_weight = weight.has_value()
? weight.value().abs() + eps
: at::ones({mean.size(0)}, mean.options());
auto inv_std = 1 / at::sqrt(var + eps);
auto scale = weight.has_value() ? abs_weight * inv_std : inv_std;
auto shift = weight.has_value() ? bias.value() - mean * abs_weight * inv_std
: -mean * inv_std;
x.mul_(normalize_shape(scale)).add_(normalize_shape(shift));
switch (activation) {
case Activation::LeakyReLU:
at::leaky_relu_(x, activation_param);
break;
case Activation::ELU:
at::elu_(x, activation_param);
break;
case Activation::Identity:
break;
}
}