wehub-resource-sync
640 lines
26 KiB
C++
640 lines
26 KiB
C++
/*
|
|
* ******************************************************************************
|
|
* *
|
|
* *
|
|
* * This program and the accompanying materials are made available under the
|
|
* * terms of the Apache License, Version 2.0 which is available at
|
|
* * https://www.apache.org/licenses/LICENSE-2.0.
|
|
* *
|
|
* * See the NOTICE file distributed with this work for additional
|
|
* * information regarding copyright ownership.
|
|
* * 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.
|
|
* *
|
|
* * SPDX-License-Identifier: Apache-2.0
|
|
* *****************************************************************************
|
|
*/
|
|
|
|
//
|
|
// @author saudet
|
|
// @author raver119@gmail.com
|
|
// @author Yurii Shyrma (iuriish@yahoo.com)
|
|
//
|
|
#include <array/NDArrayFactory.h>
|
|
#include <helpers/MKLDNNStream.h>
|
|
#include <ops/declarable/OpRegistrator.h>
|
|
#include <ops/declarable/PlatformHelper.h>
|
|
#include <ops/declarable/helpers/convolutions.h>
|
|
#include <system/platform_boilerplate.h>
|
|
|
|
#include "mkldnnUtils.h"
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace platforms {
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
static void batchnormMKLDNN(NDArray* x, NDArray* mean, NDArray* variance, NDArray* weights,
|
|
NDArray* z, const float epsilon, const bool isNCHW) {
|
|
// unfortunately mkl dnn doesn't support any format (dnnl::memory::format_tag::any) for x
|
|
|
|
// x -> 2D:nc, 4D:nchw/nhwc, 5D:ncdhw/ndhwc
|
|
// mean -> 1D [c]
|
|
// variance -> 1D [c]
|
|
// weights 2D [2, c], weights({0,1, 0,0}) contains gamma and weights({1,2, 0,0}) contains beta
|
|
// z(output) - same shape as x
|
|
|
|
const int xRank = x->rankOf();
|
|
|
|
// input type
|
|
dnnl::memory::data_type type = dnnl::memory::data_type::f32;
|
|
|
|
// indicate whether gamma or/and beta are given
|
|
auto flags =
|
|
dnnl::normalization_flags::use_global_stats; // don't calculate the mean and variance for each mini-batch
|
|
if (weights != nullptr) flags |= dnnl::normalization_flags::use_scale_shift;
|
|
|
|
dnnl::memory::dims dims;
|
|
dnnl::memory::format_tag format;
|
|
|
|
const sd::LongType indHW = isNCHW ? 2 : 1;
|
|
const sd::LongType bS = x->sizeAt(0);
|
|
const sd::LongType iC = isNCHW ? x->sizeAt(1) : x->sizeAt(-1);
|
|
|
|
int iD, iH, iW;
|
|
|
|
if (xRank == 2) {
|
|
dims = {bS, iC};
|
|
format = dnnl::memory::format_tag::nc;
|
|
} else if (xRank == 4) {
|
|
iH = x->sizeAt(indHW);
|
|
iW = x->sizeAt(indHW + 1);
|
|
dims = {bS, iC, iH, iW};
|
|
format = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
|
} else { // xRank = 5
|
|
iD = x->sizeAt(indHW);
|
|
iH = x->sizeAt(indHW + 1);
|
|
iW = x->sizeAt(indHW + 2);
|
|
dims = {bS, iC, iD, iH, iW};
|
|
format = isNCHW ? dnnl::memory::format_tag::ncdhw : dnnl::memory::format_tag::ndhwc;
|
|
}
|
|
|
|
// memory descriptors for arrays
|
|
|
|
// x
|
|
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(dims, type, format);
|
|
dnnl::memory::desc x_user_md = dnnl::memory::desc(dims, type, format);
|
|
onednnUtils::setBlockStrides(*x, x_user_md);
|
|
|
|
// z, output
|
|
dnnl::memory::desc z_mkl_md = dnnl::memory::desc(dims, type, dnnl::memory::format_tag::any);
|
|
dnnl::memory::desc z_user_md = dnnl::memory::desc(dims, type, format);
|
|
onednnUtils::setBlockStrides(*z, z_user_md);
|
|
|
|
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
|
|
|
// batchnorm forward description
|
|
dnnl::batch_normalization_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference, x_mkl_md, epsilon, flags);
|
|
dnnl::batch_normalization_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
|
|
|
// arguments (memory buffers) necessary for calculations
|
|
std::unordered_map<int, dnnl::memory> args;
|
|
|
|
dnnl::stream stream(engine);
|
|
|
|
// provide memory and check whether reorder is required
|
|
|
|
// x
|
|
onednnUtils::loadDataToMklStream(*x, engine, stream, x_user_md, op_ff_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
|
|
|
// z
|
|
auto z_user_mem =
|
|
onednnUtils::loadDataToMklStream(*z, engine, stream, z_user_md, op_ff_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
|
|
|
// mean
|
|
auto mean_mkl_mem = dnnl::memory(op_ff_prim_desc.mean_desc(), engine, const_cast<void*>(mean->buffer()));
|
|
args[DNNL_ARG_MEAN] = mean_mkl_mem;
|
|
|
|
// variance
|
|
auto var_mkl_mem = dnnl::memory(op_ff_prim_desc.variance_desc(), engine, const_cast<void*>(variance->buffer()));
|
|
args[DNNL_ARG_VARIANCE] = var_mkl_mem;
|
|
|
|
// gamma and beta (and their gradients) if they are present
|
|
if (weights != nullptr) {
|
|
auto w_mkl_mem = dnnl::memory(op_ff_prim_desc.weights_desc(), engine, const_cast<void*>(weights->buffer()));
|
|
args[DNNL_ARG_WEIGHTS] = w_mkl_mem;
|
|
}
|
|
|
|
// run calculations
|
|
dnnl::batch_normalization_forward(op_ff_prim_desc).execute(stream, args);
|
|
|
|
// reorder outputs if necessary
|
|
if (op_ff_prim_desc.dst_desc() != z_user_mem.get_desc())
|
|
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
|
|
|
stream.wait();
|
|
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
static void batchnormBpMKLDNN(NDArray* x, NDArray* mean, NDArray* variance, NDArray* dLdO,
|
|
NDArray* weights, NDArray* dLdI, NDArray* dLdW, const float epsilon,
|
|
const bool isNCHW) {
|
|
// unfortunately mkl dnn doesn't support any format (dnnl::memory::format_tag::any) for x
|
|
|
|
// x -> 2D:nc, 4D:nchw/nhwc, 5D:ncdhw/ndhwc
|
|
// mean -> 1D [c]
|
|
// variance -> 1D [c]
|
|
// dLdO - same shape as x
|
|
// weights 2D [2, c], weights({0,1, 0,0}) contains gamma and weights({1,2, 0,0}) contains beta
|
|
// dLdI - same shape as x
|
|
// dLdW - same shape as weights, dLdW({0,1, 0,0}) contains grad_gamma and dLdW({1,2, 0,0}) contains grad_beta
|
|
|
|
const sd::LongType xRank = x->rankOf();
|
|
|
|
// input type
|
|
dnnl::memory::data_type type = dnnl::memory::data_type::f32;
|
|
|
|
// indicate whether gamma or/and beta are given
|
|
auto flags =
|
|
dnnl::normalization_flags::use_global_stats; // don't calculate the mean and variance for each mini-batch
|
|
if (weights != nullptr) flags |= dnnl::normalization_flags::use_scale_shift;
|
|
|
|
dnnl::memory::dims dims;
|
|
dnnl::memory::format_tag format;
|
|
|
|
const sd::LongType indHW = isNCHW ? 2 : 1;
|
|
const sd::LongType bS = x->sizeAt(0);
|
|
const sd::LongType iC = isNCHW ? x->sizeAt(1) : x->sizeAt(-1);
|
|
|
|
sd::LongType iD, iH, iW;
|
|
|
|
if (xRank == 2) {
|
|
dims = {bS, iC};
|
|
format = dnnl::memory::format_tag::nc;
|
|
} else if (xRank == 4) {
|
|
iH = x->sizeAt(indHW);
|
|
iW = x->sizeAt(indHW + 1);
|
|
dims = {bS, iC, iH, iW};
|
|
format = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
|
} else { // xRank = 5
|
|
iD = x->sizeAt(indHW);
|
|
iH = x->sizeAt(indHW + 1);
|
|
iW = x->sizeAt(indHW + 2);
|
|
dims = {bS, iC, iD, iH, iW};
|
|
format = isNCHW ? dnnl::memory::format_tag::ncdhw : dnnl::memory::format_tag::ndhwc;
|
|
}
|
|
|
|
// memory descriptors for arrays
|
|
|
|
// x
|
|
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(dims, type, format);
|
|
dnnl::memory::desc x_user_md = dnnl::memory::desc(dims, type, format);
|
|
onednnUtils::setBlockStrides(*x, x_user_md);
|
|
|
|
// dLdO
|
|
dnnl::memory::desc dLdO_mkl_md = dnnl::memory::desc(dims, type, dnnl::memory::format_tag::any);
|
|
dnnl::memory::desc dLdO_user_md = dnnl::memory::desc(dims, type, format);
|
|
onednnUtils::setBlockStrides(*dLdO, dLdO_user_md);
|
|
|
|
// dLdI
|
|
dnnl::memory::desc dLdI_mkl_md = dnnl::memory::desc(dims, type, dnnl::memory::format_tag::any);
|
|
dnnl::memory::desc dLdI_user_md = dnnl::memory::desc(dims, type, format);
|
|
onednnUtils::setBlockStrides(*dLdI, dLdI_user_md);
|
|
|
|
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
|
|
|
// batchnorm forward description
|
|
dnnl::batch_normalization_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference, x_mkl_md, epsilon, flags);
|
|
dnnl::batch_normalization_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
|
|
|
// batchnorm backprop description
|
|
dnnl::batch_normalization_backward::desc op_bp_desc(dnnl::prop_kind::backward, dLdO_mkl_md, x_mkl_md, epsilon, flags);
|
|
dnnl::batch_normalization_backward::primitive_desc op_bp_prim_desc(op_bp_desc, engine, op_ff_prim_desc);
|
|
|
|
// arguments (memory buffers) necessary for calculations
|
|
std::unordered_map<int, dnnl::memory> args;
|
|
|
|
dnnl::stream stream(engine);
|
|
|
|
// provide memory and check whether reorder is required
|
|
|
|
// x
|
|
onednnUtils::loadDataToMklStream(*x, engine, stream, x_user_md, op_bp_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
|
|
|
// dLdO
|
|
onednnUtils::loadDataToMklStream(*dLdO, engine, stream, dLdO_user_md, op_bp_prim_desc.diff_dst_desc(),
|
|
args[DNNL_ARG_DIFF_DST]);
|
|
|
|
// mean
|
|
auto mean_mkl_mem = dnnl::memory(op_bp_prim_desc.mean_desc(), engine, const_cast<void*>(mean->buffer()));
|
|
args[DNNL_ARG_MEAN] = mean_mkl_mem;
|
|
|
|
// variance
|
|
auto var_mkl_mem = dnnl::memory(op_bp_prim_desc.variance_desc(), engine, const_cast<void*>(variance->buffer()));
|
|
args[DNNL_ARG_VARIANCE] = var_mkl_mem;
|
|
|
|
// dLdI
|
|
auto dLdI_user_mem = onednnUtils::loadDataToMklStream(*dLdI, engine, stream, dLdI_user_md,
|
|
op_bp_prim_desc.diff_src_desc(), args[DNNL_ARG_DIFF_SRC]);
|
|
|
|
// gamma and beta (and their gradients) if they are present
|
|
if (weights != nullptr) {
|
|
auto w_mkl_mem = dnnl::memory(op_bp_prim_desc.weights_desc(), engine, const_cast<void*>(weights->buffer()));
|
|
args[DNNL_ARG_WEIGHTS] = w_mkl_mem;
|
|
|
|
auto dLdW_mkl_mem = dnnl::memory(op_bp_prim_desc.weights_desc(), engine, dLdW->buffer());
|
|
args[DNNL_ARG_DIFF_WEIGHTS] = dLdW_mkl_mem;
|
|
}
|
|
|
|
// run calculations
|
|
dnnl::batch_normalization_backward(op_bp_prim_desc).execute(stream, args);
|
|
|
|
// reorder outputs if necessary
|
|
if (op_bp_prim_desc.diff_src_desc() != dLdI_user_mem.get_desc())
|
|
dnnl::reorder(args[DNNL_ARG_DIFF_SRC], dLdI_user_mem).execute(stream, args[DNNL_ARG_DIFF_SRC], dLdI_user_mem);
|
|
|
|
stream.wait();
|
|
|
|
|
|
// notations:
|
|
// f = g * (gamma * ((x - m) / (v + eps)^0.5) + beta) -> means dLdO * ff_output
|
|
// g = dLdO
|
|
// stdInv = 1 / (v + eps)^0.5
|
|
// N - batch size (product of spatial dimensions)
|
|
|
|
// formula for full derivative with respect to input (x)
|
|
// dLdI = dfdx + dfdm*dmdx + dfdv*(dvdm*dmdx + dvdx)
|
|
|
|
// !!! MKL CALCULATES ONLY FIRST TERM dfdx, SO WE SHOULD CALCULATE TERM (dfdm*dmdx + dfdv*(dvdm*dmdx + dvdx)) BY
|
|
// OURSELF !!!
|
|
|
|
// dfdm = -gamma*stdInv*g_sum;
|
|
// dmdx = 1/N;
|
|
// dvdx = 2 * (x - m) / N
|
|
// dvdm = -2 * [(x - m)]_sum / N
|
|
// dfdv = -0.5 * [g*(x - m)]_sum * stdInv^3, drop gamma here for calc convenience
|
|
|
|
// finally:
|
|
// dLdI = dfdm / N + (2/N) * dfdv * (dvdm/2 + (x - m))
|
|
// dLdI = gamma * ( stdInv * -g_sum/N + (2/N) * dfdv * (dvdm/2 + (x - m)) )
|
|
|
|
std::vector<sd::LongType> axes = isNCHW ? std::vector<sd::LongType>{1} : std::vector<sd::LongType>{xRank - 1};
|
|
const auto excludedAxes = ShapeUtils::evalDimsToExclude(x->rankOf(),axes.size(), axes.data());
|
|
|
|
// inversed batch size 1 / N
|
|
const auto Ninv = 1.f * mean->lengthOf() / x->lengthOf();
|
|
|
|
// x - mean
|
|
NDArray xMinusMean(x); // empty array with same shape as x
|
|
const_cast<NDArray*>(x)->applyBroadcast(sd::broadcast::Subtract, &axes, mean, &xMinusMean);
|
|
|
|
// stdInv
|
|
NDArray stdInv = *variance + epsilon;
|
|
stdInv.applyTransform(transform::Reciprocal, &stdInv); // 1 / (variance + epsilon)
|
|
stdInv.applyTransform(transform::Sqrt, &stdInv); // 1 / (variance + epsilon)^0.5
|
|
|
|
// dfdm / N
|
|
auto dfdm = dLdO->reduceAlongDimension(sd::reduce::Sum, excludedAxes);
|
|
dfdm *= stdInv;
|
|
dfdm *= -Ninv;
|
|
|
|
// dvdm / 2
|
|
NDArray dvdm(mean); // empty array with same shape as mean
|
|
xMinusMean.reduceAlongDimension(sd::reduce::Sum, &dvdm, excludedAxes);
|
|
dvdm *= -Ninv;
|
|
|
|
// (2/N)*dfdv
|
|
NDArray dfdv(variance); // empty array with same shape as variance
|
|
(xMinusMean * *dLdO).reduceAlongDimension(sd::reduce::Sum, &dfdv, excludedAxes);
|
|
dfdv *= stdInv * stdInv * stdInv;
|
|
dfdv *= -Ninv;
|
|
|
|
// dvdm/2 + (x - m)
|
|
xMinusMean.applyBroadcast(sd::broadcast::Add, &axes, &dvdm, &xMinusMean);
|
|
// dfdv * (dvdm/2 + (x - m))
|
|
xMinusMean.applyBroadcast(sd::broadcast::Multiply, &axes, &dfdv, &xMinusMean);
|
|
// add dfdm / N
|
|
xMinusMean.applyBroadcast(sd::broadcast::Add, &axes, &dfdm, &xMinusMean);
|
|
// * gamma
|
|
auto gamma = (*weights)({0, 1, 0, 0});
|
|
xMinusMean.applyBroadcast(sd::broadcast::Multiply, &axes, &gamma, &xMinusMean);
|
|
|
|
*dLdI += xMinusMean;
|
|
}
|
|
|
|
PLATFORM_IMPL(batchnorm, ENGINE_CPU) {
|
|
auto input = INPUT_VARIABLE(0); // 2D:nc, 4D:nchw/nhwc, 5D:ncdhw/ndhwc
|
|
auto mean = INPUT_VARIABLE(1); // [c]
|
|
auto variance = INPUT_VARIABLE(2); // [c]
|
|
NDArray* gamma = nullptr; // [c]
|
|
NDArray* beta = nullptr; // [c]
|
|
|
|
auto output = OUTPUT_VARIABLE(0); // same shape as input
|
|
|
|
const bool applyScale = (bool)INT_ARG(0);
|
|
const bool applyOffset = (bool)INT_ARG(1);
|
|
const double epsilon = T_ARG(0);
|
|
|
|
if (applyScale) gamma = INPUT_VARIABLE(3);
|
|
if (applyOffset) beta = INPUT_VARIABLE(3 + (int)applyScale);
|
|
|
|
const int numOfIntArgs = block.getIArguments()->size();
|
|
const sd::LongType inRank = input->rankOf();
|
|
|
|
// get axes args to normalize input array over
|
|
std::vector<sd::LongType> axes;
|
|
if (numOfIntArgs > 2)
|
|
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
|
|
else
|
|
axes.push_back(inRank - 1); // default dimension to reduce along is last dimension
|
|
|
|
const sd::LongType numOfAxes = axes.size();
|
|
REQUIRE_TRUE(numOfAxes == 1, 0,
|
|
"BATCHNORM_MKLDNN op: mkl dnn library supports only one axis which represents channel dimension, but "
|
|
"got %i axes instead!",
|
|
numOfAxes);
|
|
REQUIRE_TRUE(inRank == 2 || inRank == 4 || inRank == 5, 0,
|
|
"BATCHNORM_MKLDNN op: possible values for rank of input array are 2, 4 or 5, but got %i instead!",
|
|
inRank);
|
|
REQUIRE_TRUE(mean->rankOf() == 1 && mean->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
|
"BATCHNORM_MKLDNN op: wrong shape of mean array, expected is [%lld], but got %s instead !",
|
|
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(mean).c_str());
|
|
REQUIRE_TRUE(variance->rankOf() == 1 && variance->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
|
"BATCHNORM_MKLDNN op: wrong shape of variance array, expected is [%lld], but got %s instead !",
|
|
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(variance).c_str());
|
|
if (gamma != nullptr)
|
|
REQUIRE_TRUE(gamma->rankOf() == 1 && gamma->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
|
"BATCHNORM_MKLDNN op: wrong shape of gamma array, expected is [%lld], but got %s instead !",
|
|
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(gamma).c_str());
|
|
if (beta != nullptr)
|
|
REQUIRE_TRUE(beta->rankOf() == 1 && beta->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
|
"BATCHNORM_MKLDNN op: wrong shape of beta array, expected is [%lld], but got %s instead !",
|
|
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(beta).c_str());
|
|
|
|
// types of all input arrays should be the same (except dLdO)
|
|
for (size_t i = 1; i < block.width() - 1; ++i)
|
|
REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
|
|
"BATCHNORM_MKLDNN op: types of all input arrays should be the same !");
|
|
|
|
NDArray* weights = nullptr;
|
|
|
|
if (applyScale || applyOffset) {
|
|
std::vector<sd::LongType > shape = {2, input->sizeAt(axes[0])};
|
|
weights = new NDArray(input->ordering(),shape , input->dataType());
|
|
|
|
if (applyScale)
|
|
(*weights)({0, 1, 0, 0}).assign(gamma);
|
|
else {
|
|
sd::LongType scalarVal = 1;
|
|
(*weights)({0, 1, 0, 0}).assign(scalarVal);
|
|
}
|
|
if (applyOffset)
|
|
(*weights)({1, 2, 0, 0}).assign(beta);
|
|
else
|
|
(*weights)({1, 2, 0, 0}).assign(0);
|
|
}
|
|
|
|
const bool isNCHW = !(axes[0] == inRank - 1 && inRank > 2);
|
|
|
|
batchnormMKLDNN(input, mean, variance, weights, output, epsilon, isNCHW);
|
|
|
|
delete weights;
|
|
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
PLATFORM_CHECK(batchnorm, ENGINE_CPU) {
|
|
auto input = INPUT_VARIABLE(0); // 2D:nc, 4D:nchw/nhwc, 5D:ncdhw/ndhwc
|
|
auto mean = INPUT_VARIABLE(1); // [c]
|
|
auto variance = INPUT_VARIABLE(2); // [c]
|
|
NDArray* gamma = nullptr; // [c]
|
|
NDArray* beta = nullptr; // [c]
|
|
|
|
auto output = OUTPUT_VARIABLE(0); // same shape as input
|
|
|
|
const bool applyScale = (bool)INT_ARG(0);
|
|
const bool applyOffset = (bool)INT_ARG(1);
|
|
|
|
if (applyScale) gamma = INPUT_VARIABLE(3);
|
|
if (applyOffset) beta = INPUT_VARIABLE(3 + (int)applyScale);
|
|
|
|
const int numOfIntArgs = block.getIArguments()->size();
|
|
std::vector<sd::LongType> axes;
|
|
if (numOfIntArgs > 2)
|
|
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
|
|
else
|
|
axes.push_back(input->rankOf() - 1); // default dimension to reduce along is last dimension
|
|
|
|
const int inRank = input->rankOf();
|
|
|
|
Requirements req("ONEDNN BATCHNORM OP");
|
|
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
|
req.expectEq(makeInfoVariable(axes.size(), "axes.size()"), 1) &&
|
|
req.expectIn(makeInfoVariable(axes[0], "axes#0"), {1, inRank - 1}) &&
|
|
req.expectIn(makeInfoVariable(inRank, RANK_MSG_INPUT0), {2, 4, 5}) &&
|
|
req.expectTrue(makeInfoVariable(
|
|
[input, mean, variance, gamma, beta, output] {
|
|
DataType inputType = input->dataType();
|
|
DataType meanType = mean->dataType();
|
|
DataType varType = variance->dataType();
|
|
DataType gammaType = gamma != nullptr ? gamma->dataType() : DataType::FLOAT32;
|
|
DataType betaType = beta != nullptr ? beta->dataType() : DataType::FLOAT32;
|
|
DataType outType = output->dataType();
|
|
return (inputType == DataType::FLOAT32 && meanType == DataType::FLOAT32 &&
|
|
varType == DataType::FLOAT32 && gammaType == DataType::FLOAT32 &&
|
|
betaType == DataType::FLOAT32 && outType == DataType::FLOAT32);
|
|
},
|
|
TYPECHECK_MSG),
|
|
NO_MSG);
|
|
req.logTheSuccess();
|
|
return req;
|
|
}
|
|
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
PLATFORM_IMPL(batchnorm_bp, ENGINE_CPU) {
|
|
NDArray* input = INPUT_VARIABLE(0); // 2D:nc, 4D:nchw/nhwc, 5D:ncdhw/ndhwc
|
|
NDArray* mean = INPUT_VARIABLE(1); // [c]
|
|
NDArray* variance = INPUT_VARIABLE(2); // [c]
|
|
NDArray* gamma = nullptr; // [c]
|
|
NDArray* beta = nullptr; // [c]
|
|
NDArray* dLdO = INPUT_VARIABLE(block.width() - 1); // same as input
|
|
|
|
NDArray* dLdI = OUTPUT_VARIABLE(0); // same as input
|
|
NDArray* dLdM = OUTPUT_VARIABLE(1); // [c]
|
|
NDArray* dLdV = OUTPUT_VARIABLE(2); // [c]
|
|
NDArray* dLdG = nullptr; // [c]
|
|
NDArray* dLdB = nullptr; // [c]
|
|
|
|
const bool applyScale = (bool)INT_ARG(0);
|
|
const bool applyOffset = (bool)INT_ARG(1);
|
|
const float epsilon = T_ARG(0);
|
|
|
|
if (applyScale) {
|
|
gamma = INPUT_VARIABLE(3);
|
|
dLdG = OUTPUT_VARIABLE(3);
|
|
}
|
|
if (applyOffset) {
|
|
beta = INPUT_VARIABLE(3 + (sd::LongType)applyScale);
|
|
dLdB = OUTPUT_VARIABLE(3 + (sd::LongType)applyScale);
|
|
}
|
|
|
|
const int numOfIntArgs = block.getIArguments()->size();
|
|
const int inRank = input->rankOf();
|
|
|
|
// get axes args to normalize input array over
|
|
std::vector<sd::LongType> axes;
|
|
if (numOfIntArgs > 2)
|
|
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
|
|
else
|
|
axes.push_back(inRank - 1); // default dimension to reduce along is last dimension
|
|
|
|
const int numOfAxes = axes.size();
|
|
REQUIRE_TRUE(numOfAxes == 1, 0,
|
|
"BATCHNORM_BP_MKLDNN op: mkl dnn library supports only one axis which represents channel dimension, but "
|
|
"got %i axes instead!",
|
|
numOfAxes);
|
|
REQUIRE_TRUE(inRank == 2 || inRank == 4 || inRank == 5, 0,
|
|
"BATCHNORM_BP_MKLDNN op: possible values for rank of input array are 2, 4 or 5, but got %i instead!",
|
|
inRank);
|
|
REQUIRE_TRUE(input->isSameShape(dLdO), 0,
|
|
"BATCHNORM_BP_MKLDNN op: wrong shape of gradients array, expected is %s, but got %s instead !",
|
|
ShapeUtils::shapeAsString(input).c_str(), ShapeUtils::shapeAsString(dLdO).c_str());
|
|
REQUIRE_TRUE(mean->rankOf() == 1 && mean->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
|
"BATCHNORM_BP_MKLDNN op: wrong shape of mean array, expected is [%lld], but got %s instead !",
|
|
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(mean).c_str());
|
|
REQUIRE_TRUE(variance->rankOf() == 1 && variance->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
|
"BATCHNORM_BP_MKLDNN op: wrong shape of variance array, expected is [%lld], but got %s instead !",
|
|
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(variance).c_str());
|
|
if (gamma != nullptr)
|
|
REQUIRE_TRUE(gamma->rankOf() == 1 && gamma->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
|
"BATCHNORM_BP_MKLDNN op: wrong shape of gamma array, expected is [%lld], but got %s instead !",
|
|
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(gamma).c_str());
|
|
if (beta != nullptr)
|
|
REQUIRE_TRUE(beta->rankOf() == 1 && beta->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
|
"BATCHNORM_BP_MKLDNN op: wrong shape of beta array, expected is [%lld], but got %s instead !",
|
|
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(beta).c_str());
|
|
|
|
// types of all input arrays should be the same
|
|
for (size_t i = 1; i < block.width() - 1; ++i)
|
|
REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
|
|
"BATCHNORM_BP_MKLDNN op: types of all input arrays should be the same !");
|
|
|
|
NDArray *weights = nullptr, *dLdW = nullptr;
|
|
|
|
if (applyScale || applyOffset) {
|
|
sd::LongType scalar = 1;
|
|
std::vector<sd::LongType> shape = {2, input->sizeAt(axes[0])};
|
|
weights = new NDArray(input->ordering(),shape, input->dataType());
|
|
dLdW = new NDArray(input->ordering(), shape, input->dataType());
|
|
if (applyScale)
|
|
(*weights)({0, 1, 0, 0}).assign(gamma);
|
|
else
|
|
(*weights)({0, 1, 0, 0}).assign(scalar);
|
|
if (applyOffset)
|
|
(*weights)({1, 2, 0, 0}).assign(beta);
|
|
else
|
|
(*weights)({1, 2, 0, 0}).assign(0);
|
|
}
|
|
|
|
const bool isNCHW = !(axes[0] == inRank - 1 && inRank > 2);
|
|
|
|
if (shape::strideDescendingCAscendingF(dLdO->shapeInfo()))
|
|
batchnormBpMKLDNN(input, mean, variance, dLdO, weights, dLdI, dLdW, epsilon, isNCHW);
|
|
else {
|
|
NDArray dupped = dLdO->dup();
|
|
batchnormBpMKLDNN(input, mean, variance, &dupped, weights, dLdI, dLdW, epsilon, isNCHW);
|
|
}
|
|
*dLdM = 0;
|
|
*dLdV = 0;
|
|
|
|
if (applyScale || applyOffset) {
|
|
if (applyScale) {
|
|
NDArray assign = (*dLdW)({0, 1, 0, 0});
|
|
dLdG->assign(&assign);
|
|
}
|
|
if (applyOffset) {
|
|
NDArray assign = (*dLdW)({1, 2, 0, 0});
|
|
dLdB->assign(&assign);
|
|
}
|
|
|
|
delete weights;
|
|
delete dLdW;
|
|
}
|
|
|
|
return sd::Status::OK;
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
PLATFORM_CHECK(batchnorm_bp, ENGINE_CPU) {
|
|
NDArray* input = INPUT_VARIABLE(0); // 2D:nc, 4D:nchw, 5D:ncdhw
|
|
NDArray* mean = INPUT_VARIABLE(1); // [c]
|
|
NDArray* variance = INPUT_VARIABLE(2); // [c]
|
|
NDArray* dLdO = INPUT_VARIABLE(3); // same as input
|
|
NDArray* gamma = nullptr; // [c]
|
|
NDArray* beta = nullptr; // [c]
|
|
|
|
NDArray* dLdI = OUTPUT_VARIABLE(0); // same as input
|
|
NDArray* dLdM = OUTPUT_VARIABLE(1); // [c]
|
|
NDArray* dLdV = OUTPUT_VARIABLE(2); // [c]
|
|
NDArray* dLdG = nullptr; // [c]
|
|
NDArray* dLdB = nullptr; // [c]
|
|
|
|
const bool applyScale = (bool)INT_ARG(0);
|
|
const bool applyOffset = (bool)INT_ARG(1);
|
|
|
|
if (applyScale) {
|
|
gamma = INPUT_VARIABLE(4);
|
|
dLdG = OUTPUT_VARIABLE(3);
|
|
}
|
|
if (applyOffset) {
|
|
beta = INPUT_VARIABLE(4 + (sd::LongType)applyScale);
|
|
dLdB = OUTPUT_VARIABLE(3 + (sd::LongType)applyScale);
|
|
}
|
|
|
|
const int numOfIntArgs = block.getIArguments()->size();
|
|
std::vector<sd::LongType> axes;
|
|
if (numOfIntArgs > 2)
|
|
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
|
|
else
|
|
axes.push_back(input->rankOf() - 1); // default dimension to reduce along is last dimension
|
|
|
|
const sd::LongType inRank = input->rankOf();
|
|
std::vector<sd::LongType> shape = {1, inRank - 1};
|
|
Requirements req("ONEDNN BATCHNORM_BP OP");
|
|
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
|
req.expectEq(makeInfoVariable(axes.size(), "axes.size()"), 1) &&
|
|
req.expectIn(makeInfoVariable(inRank, RANK_MSG_INPUT0), {2, 4, 5}) &&
|
|
req.expectTrue(makeInfoVariable(
|
|
[input, mean, variance, dLdO, gamma, beta, dLdG, dLdB, dLdI] {
|
|
DataType inputType = input->dataType();
|
|
DataType meanType = mean->dataType();
|
|
DataType varType = variance->dataType();
|
|
DataType dLdOType = dLdO->dataType();
|
|
DataType gammaType = gamma != nullptr ? gamma->dataType() : DataType::FLOAT32;
|
|
DataType betaType = beta != nullptr ? beta->dataType() : DataType::FLOAT32;
|
|
|
|
DataType dLdIType = dLdI->dataType();
|
|
DataType dLdGType = gamma != nullptr ? dLdG->dataType() : DataType::FLOAT32;
|
|
DataType dLdBType = beta != nullptr ? dLdB->dataType() : DataType::FLOAT32;
|
|
return (inputType == DataType::FLOAT32 && meanType == DataType::FLOAT32 &&
|
|
varType == DataType::FLOAT32 && dLdOType == DataType::FLOAT32 &&
|
|
gammaType == DataType::FLOAT32 && betaType == DataType::FLOAT32 &&
|
|
dLdIType == DataType::FLOAT32 && dLdGType == DataType::FLOAT32 &&
|
|
dLdBType == DataType::FLOAT32);
|
|
},
|
|
TYPECHECK_MSG),
|
|
NO_MSG);
|
|
req.logTheSuccess();
|
|
return req;
|
|
}
|
|
|
|
} // namespace platforms
|
|
} // namespace ops
|
|
} // namespace sd
|