diff --git a/modules/dnn/src/graph_fusion_basic.cpp b/modules/dnn/src/graph_fusion_basic.cpp index 6a8cc52daf..2d6bd21c60 100644 --- a/modules/dnn/src/graph_fusion_basic.cpp +++ b/modules/dnn/src/graph_fusion_basic.cpp @@ -280,5 +280,179 @@ void Net::Impl::fuseBasic() basicFusion.fuse(); } +// fold BN scale/bias into the weights of the immediately following Conv2 (pre-constArgs only) +struct FuseBNPass +{ + FuseBNPass(Net::Impl* netimpl_) : netimpl(netimpl_) {} + + void run() + { + netimpl->useCounts(usecounts); + fuseGraph(netimpl->mainGraph); + } + + void fuseGraph(Ptr& graph) + { + const std::vector >& prog = graph->prog(); + size_t nops = prog.size(), nargs = netimpl->args.size(); + std::vector > newprog; + newprog.reserve(nops); + std::vector producer_of((int)nargs, -1); + bool modified = false; + + for (size_t i = 0; i < nops; i++) { + const Ptr& layer = prog[i]; + Layer* layer_ptr = const_cast(layer.get()); + + std::vector >* subgraphs = layer->subgraphs(); + if (subgraphs) + for (Ptr& g : *subgraphs) fuseGraph(g); + + const std::vector& inputs = layer->inputs; + const std::vector& outputs = layer->outputs; + + Conv2Layer* conv = dynamic_cast(layer_ptr); + if (conv && !inputs.empty()) { + Arg conv_inp0 = inputs[0]; + int bn_idx = conv_inp0.idx >= 0 && conv_inp0.idx < (int)producer_of.size() + ? producer_of[conv_inp0.idx] : -1; + if (bn_idx >= 0 && usecounts[conv_inp0.idx] == 1) { + BatchNorm2Layer* bn = dynamic_cast(newprog[bn_idx].get()); + if (bn && fuseForward(conv, bn)) { + Arg bn_inp0 = bn->inputs[0]; + layer_ptr->inputs[0] = bn_inp0; + usecounts[conv_inp0.idx] = 0; + if (bn_inp0.idx >= 0) + usecounts[bn_inp0.idx]++; + newprog[bn_idx] = Ptr(); + modified = true; + } + } + } + + for (Arg out : outputs) + if (out.idx >= 0 && out.idx < (int)producer_of.size()) + producer_of[out.idx] = (int)newprog.size(); + newprog.push_back(layer); + } + + if (modified) { + size_t j = 0; + for (size_t i = 0; i < newprog.size(); i++) { + if (!newprog[i].empty()) { + if (j < i) newprog[j] = newprog[i]; + j++; + } + } + newprog.resize(j); + graph->setProg(newprog); + } + } + + // fold BN scale/bias into Conv2 weight and bias tensors (raw NCHW fp32 only) + bool fuseForward(Conv2Layer* conv, BatchNorm2Layer* bn) + { + // skip padded convolutions: border taps get incorrect bias from folded BN + for (int p : conv->pads) + if (p > 0) return false; + + if (conv->inputs.size() < 2 || !netimpl->isConstArg(conv->inputs[1])) + return false; + + Mat& Wref = netimpl->argTensor(conv->inputs[1]); + if (Wref.empty() || Wref.type() != CV_32F || Wref.dims != 4) + return false; + + int OC = Wref.size[0]; + int IC_pg = Wref.size[1]; // input channels per group + int KH = Wref.size[2]; + int KW = Wref.size[3]; + int ngroups = conv->ngroups; + if (ngroups <= 0 || OC % ngroups != 0) + return false; + int OC_pg = OC / ngroups; + + // resolve BN effective scale and bias; validate all params are const + Mat bn_scale, bn_bias_vec; + size_t bn_nin = bn->inputs.size(); + if (bn_nin == 5) { + for (int k = 1; k <= 4; k++) + if (!netimpl->isConstArg(bn->inputs[k])) return false; + BatchNorm2Layer::getScaleBias( + netimpl->argTensor(bn->inputs[1]), + netimpl->argTensor(bn->inputs[2]), + netimpl->argTensor(bn->inputs[3]), + netimpl->argTensor(bn->inputs[4]), + bn->epsilon, bn_scale, bn_bias_vec); + } else if (bn_nin == 1) { + bn->getScaleBias(bn_scale, bn_bias_vec); + } else { + return false; + } + + bn_scale.convertTo(bn_scale, CV_32F); + bn_bias_vec.convertTo(bn_bias_vec, CV_32F); + + if ((int)bn_scale.total() != IC_pg * ngroups) + return false; + + const float* bn_s = bn_scale.ptr(); + const float* bn_b = bn_bias_vec.ptr(); + + // reshape to 2D [OC, IC_pg*KH*KW]; shares buffer with Wref + Mat W2d = Wref.reshape(1, OC); + + // scale weights in-place; accumulate per-output-channel bias correction + bias_adj.assign(OC, 0.f); + for (int g = 0; g < ngroups; g++) { + int ic_off = g * IC_pg; + for (int oc = g * OC_pg, oc_end = oc + OC_pg; oc < oc_end; oc++) { + float* wrow = W2d.ptr(oc); + for (int ic = 0; ic < IC_pg; ic++) { + float s = bn_s[ic_off + ic]; + float b = bn_b[ic_off + ic]; + float* wic = wrow + ic * KH * KW; + float sw = 0.f; + for (int k = 0; k < KH * KW; k++) { + sw += wic[k]; + wic[k] *= s; + } + bias_adj[oc] += b * sw; + } + } + } + + // apply bias correction into existing bias tensor or create a new one + size_t conv_nin = conv->inputs.size(); + if (conv_nin >= 3 && netimpl->isConstArg(conv->inputs[2])) { + Mat& Bref = netimpl->argTensor(conv->inputs[2]); + if (Bref.type() != CV_32F || (int)Bref.total() != OC) + return false; + float* bp = Bref.ptr(); + for (int oc = 0; oc < OC; oc++) + bp[oc] += bias_adj[oc]; + } else { + Mat newB(1, OC, CV_32F); + std::memcpy(newB.ptr(), bias_adj.data(), OC * sizeof(float)); + Arg ba = netimpl->newConstArg( + "__fused_bn_bias_w" + std::to_string(conv->inputs[1].idx), newB); + if (conv_nin == 2) conv->inputs.push_back(ba); + else conv->inputs[2] = ba; + } + + return true; + } + + Net::Impl* netimpl; + std::vector usecounts; + std::vector bias_adj; +}; + +void Net::Impl::fuseBN() +{ + FuseBNPass pass(this); + pass.run(); +} + CV__DNN_INLINE_NS_END }} diff --git a/modules/dnn/src/layers/conv2_layer.cpp b/modules/dnn/src/layers/conv2_layer.cpp index cde386b5f2..d6d46a1d84 100644 --- a/modules/dnn/src/layers/conv2_layer.cpp +++ b/modules/dnn/src/layers/conv2_layer.cpp @@ -157,7 +157,8 @@ public: virtual bool fuseBatchNorm(const Ptr& bnlayer) override { BatchNorm2Layer* bn = dynamic_cast(bnlayer.get()); - if (fusedBatchNorm || !bn || bn->inputs.size() > 1) + if (fusedBatchNorm || !bn || bn->inputs.size() > 1 || + fastActivation != FAST_ACTIV_NONE || !activ.empty()) return false; fuseBatchNormWeights(bn); fusedBatchNorm = true; diff --git a/modules/dnn/src/net_impl.hpp b/modules/dnn/src/net_impl.hpp index a7b453ebbe..832be12cf1 100644 --- a/modules/dnn/src/net_impl.hpp +++ b/modules/dnn/src/net_impl.hpp @@ -474,6 +474,8 @@ struct Net::Impl : public detail::NetImplBase // insert transformLayout operations where necessary; // use block layout for convolution, pooling and some other operations where it matters void useBlockLayout(); + // fuse BN into following Conv2 weights + void fuseBN(); }; // Net::Impl diff --git a/modules/dnn/src/net_impl2.cpp b/modules/dnn/src/net_impl2.cpp index 69ad9830bf..42bce5e8f8 100644 --- a/modules/dnn/src/net_impl2.cpp +++ b/modules/dnn/src/net_impl2.cpp @@ -547,6 +547,7 @@ void Net::Impl::prepareForInference() if (!prepared) { fuseQDQ(); constFold(); + fuseBN(); constArgs(); fuseAttention(); fuseSharedInputGemm(); diff --git a/modules/dnn/test/test_onnx_importer.cpp b/modules/dnn/test/test_onnx_importer.cpp index 7c227942b9..57ba14794a 100644 --- a/modules/dnn/test/test_onnx_importer.cpp +++ b/modules/dnn/test/test_onnx_importer.cpp @@ -943,6 +943,38 @@ TEST_P(Test_ONNX_layers, BatchNormalization) testONNXModels("batch_norm"); } +TEST_P(Test_ONNX_layers, Colorization) +{ + applyTestTag( + target == DNN_TARGET_CPU ? CV_TEST_TAG_MEMORY_512MB : CV_TEST_TAG_MEMORY_1GB, + CV_TEST_TAG_DEBUG_VERYLONG + ); + checkBackend(); + + const String model = findDataFile("dnn/colorization_deploy_v2.onnx", false); + Net net = readNetFromONNX(model); + ASSERT_FALSE(net.empty()); + net.setPreferableBackend(backend); + net.setPreferableTarget(target); + + if (target == DNN_TARGET_CPU_FP16) + net.enableWinograd(false); + + Mat inp = blobFromNPY(findDataFile("dnn/colorization_inp.npy")); + Mat ref = blobFromNPY(findDataFile("dnn/colorization_out.npy")); + + net.setInput(inp); + Mat out = net.forward(); + + double l1 = 4e-4, lInf = 3e-3; + if (target == DNN_TARGET_OPENCL_FP16 || target == DNN_TARGET_CPU_FP16 || target == DNN_TARGET_CUDA_FP16) + { + l1 = 0.25; + lInf = 5.4; + } + normAssert(out, ref, "", l1, lInf); +} + TEST_P(Test_ONNX_layers, BatchNormalization3D) { if (backend == DNN_BACKEND_INFERENCE_ENGINE_NN_BUILDER_2019)