diff --git a/modules/dnn/src/layers/batch_norm_layer.cpp b/modules/dnn/src/layers/batch_norm_layer.cpp index 4c69c247c4..b2fd75aef1 100644 --- a/modules/dnn/src/layers/batch_norm_layer.cpp +++ b/modules/dnn/src/layers/batch_norm_layer.cpp @@ -29,6 +29,8 @@ class BatchNormLayerImpl CV_FINAL : public BatchNormLayer public: Mat weights_, bias_; UMat umat_weight, umat_bias; + mutable int dims; + BatchNormLayerImpl(const LayerParams& params) { @@ -142,6 +144,7 @@ public: std::vector &outputs, std::vector &internals) const CV_OVERRIDE { + dims = inputs[0].size(); if (!useGlobalStats && inputs[0][0] != 1) CV_Error(Error::StsNotImplemented, "Batch normalization in training mode with batch size > 1"); Layer::getMemoryShapes(inputs, requiredOutputs, outputs, internals); @@ -150,9 +153,9 @@ public: virtual bool supportBackend(int backendId) CV_OVERRIDE { - return backendId == DNN_BACKEND_OPENCV || + return (backendId == DNN_BACKEND_OPENCV) || (backendId == DNN_BACKEND_HALIDE && haveHalide()) || - (backendId == DNN_BACKEND_INFERENCE_ENGINE && haveInfEngine()); + (backendId == DNN_BACKEND_INFERENCE_ENGINE && haveInfEngine() && (preferableTarget == DNN_TARGET_CPU || dims == 4)); } #ifdef HAVE_OPENCL @@ -178,11 +181,12 @@ public: } UMat &inpBlob = inputs[0]; - CV_Assert(inpBlob.dims == 2 || inpBlob.dims == 4); int groups = inpBlob.size[0]; int channels = inpBlob.size[1]; - int rows = inpBlob.dims > 2 ? inpBlob.size[2] : 1; - int cols = inpBlob.dims > 2 ? inpBlob.size[3] : 1; + int planeSize = 1; + for (size_t i = 2; i < inpBlob.dims; i++) { + planeSize *= inpBlob.size[i]; + } String opts = (use_half) ? " -DDtype=half" : " -DDtype=float"; for (size_t ii = 0; ii < outputs.size(); ii++) @@ -196,7 +200,7 @@ public: } else { - MatShape s = shape(groups * channels, rows * cols); + MatShape s = shape(groups * channels, planeSize); UMat src = inputs[ii].reshape(1, s.size(), &s[0]); UMat dst = outputs[ii].reshape(1, s.size(), &s[0]); int number = (s[1] % 8 == 0) ? 8 : ((s[1] % 4 == 0) ? 4 : 1); @@ -248,9 +252,10 @@ public: CV_Assert(inputs.size() == 1); Mat &inpBlob = inputs[0]; - CV_Assert(inpBlob.dims == 2 || inpBlob.dims == 4); - int rows = inpBlob.dims > 2 ? inpBlob.size[2] : 1; - int cols = inpBlob.dims > 2 ? inpBlob.size[3] : 1; + int planeSize = 1; + for (size_t i = 2; i < inpBlob.dims; i++) { + planeSize *= inpBlob.size[i]; + } for (size_t ii = 0; ii < outputs.size(); ii++) { @@ -262,8 +267,8 @@ public: { float w = weights_.at(n); float b = bias_.at(n); - Mat inpBlobPlane(rows, cols, CV_32F, inpBlob.ptr(num, n)); - Mat outBlobPlane(rows, cols, CV_32F, outBlob.ptr(num, n)); + Mat inpBlobPlane(1, planeSize, CV_32F, inpBlob.ptr(num, n)); + Mat outBlobPlane(1, planeSize, CV_32F, outBlob.ptr(num, n)); inpBlobPlane.convertTo(outBlobPlane, CV_32F, w, b); } } diff --git a/modules/dnn/test/test_onnx_importer.cpp b/modules/dnn/test/test_onnx_importer.cpp index f1b0a81e8e..bf9f25d214 100644 --- a/modules/dnn/test/test_onnx_importer.cpp +++ b/modules/dnn/test/test_onnx_importer.cpp @@ -167,6 +167,13 @@ TEST_P(Test_ONNX_layers, BatchNormalization) testONNXModels("batch_norm"); } +TEST_P(Test_ONNX_layers, BatchNormalization3D) +{ + if (backend == DNN_BACKEND_INFERENCE_ENGINE && target != DNN_TARGET_CPU) + throw SkipTestException(""); + testONNXModels("batch_norm_3d"); +} + TEST_P(Test_ONNX_layers, Transpose) { if (backend == DNN_BACKEND_INFERENCE_ENGINE && diff --git a/modules/dnn/test/test_tf_importer.cpp b/modules/dnn/test/test_tf_importer.cpp index e662da53e7..2c7ae44783 100644 --- a/modules/dnn/test/test_tf_importer.cpp +++ b/modules/dnn/test/test_tf_importer.cpp @@ -188,6 +188,13 @@ TEST_P(Test_TensorFlow_layers, batch_norm) runTensorFlowNet("mvn_batch_norm_1x1"); } +TEST_P(Test_TensorFlow_layers, batch_norm3D) +{ + if (backend == DNN_BACKEND_INFERENCE_ENGINE && target != DNN_TARGET_CPU) + throw SkipTestException(""); + runTensorFlowNet("batch_norm3d"); +} + TEST_P(Test_TensorFlow_layers, slim_batch_norm) { if (backend == DNN_BACKEND_INFERENCE_ENGINE)