From 34600510e4d731d057e606531f2c78387a8b8b3b Mon Sep 17 00:00:00 2001 From: Abhishek Gola Date: Wed, 5 Aug 2026 18:48:33 +0530 Subject: [PATCH] lstm optimizations --- modules/dnn/src/layers/recurrent2_layers.cpp | 505 +++++++++++-------- modules/dnn/src/onnx/onnx_importer2.cpp | 8 + 2 files changed, 297 insertions(+), 216 deletions(-) diff --git a/modules/dnn/src/layers/recurrent2_layers.cpp b/modules/dnn/src/layers/recurrent2_layers.cpp index 36cdd4ebef..6eb149fe22 100644 --- a/modules/dnn/src/layers/recurrent2_layers.cpp +++ b/modules/dnn/src/layers/recurrent2_layers.cpp @@ -3,10 +3,10 @@ // of this distribution and at http://opencv.org/license.html. #include "../precomp.hpp" -#include #include #include #include "layers_common.hpp" +#include "cpu_kernels/fast_gemm.hpp" #include "../net_impl.hpp" namespace cv @@ -41,6 +41,24 @@ static void sigmoid(const Mat &src, Mat &dst) cv::pow(1 + dst, -1, dst); } +// ONNX gate order is I,O,F,C but the recurrence wants I,F,O,C: swap the O/F blocks. +template +static void reorderGatesIOFCtoIFOC_(Mat& in) +{ + int first = in.size[0]; + int rest = in.total() / first / 4; + Mat m = in.reshape(1, {first, 4, rest}); + Mat outputGate = m.col(1), forgetGate = m.col(2); + std::swap_ranges(outputGate.begin(), outputGate.end(), forgetGate.begin()); +} +static void reorderGatesIOFCtoIFOC(Mat& in) +{ + if (in.depth() == CV_64F) + reorderGatesIOFCtoIFOC_(in); + else + reorderGatesIOFCtoIFOC_(in); +} + typedef void (*ActivationFunction)(const Mat &src, Mat &dst); static ActivationFunction get_activation_function(const String& activation) { if (activation == "Tanh"){ @@ -78,6 +96,13 @@ class LSTM2LayerImpl CV_FINAL : public LSTM2Layer ActivationFunction g_activation; ActivationFunction h_activation; + bool useAVX, useAVX2, useSVE, useNEON; + bool constWeights; // W/R/B are graph constants -> transform them once + bool weightsPacked; // whether the cached transform has been built + std::vector weightBlobs; // cached transformed weights: [Wx, Wh, bias, (pI,pF,pO)] + FastGemmOpt gemmOpt; // MLAS GEMM options for the batched input projection + std::vector> packedWx; // per-direction prepacked Wx (fp32) for fastGemm + public: LSTM2LayerImpl(const LayerParams& params) { @@ -112,6 +137,14 @@ class LSTM2LayerImpl CV_FINAL : public LSTM2Layer h_activation = get_activation_function(activations.getStringValue(2)); } + constWeights = params.get("const_weights", false); + weightsPacked = false; + useAVX = checkHardwareSupport(CPU_AVX); + useAVX2 = checkHardwareSupport(CPU_AVX2); + useSVE = checkHardwareSupport(CPU_SVE); + useNEON = checkHardwareSupport(CPU_NEON); + gemmOpt.init(); + outTailShape.clear(); } @@ -184,11 +217,8 @@ class LSTM2LayerImpl CV_FINAL : public LSTM2Layer outputs.push_back(newShape); } - // internal shapes need during forward pass - internals.assign(1, shape(_batchSize, _hidSize)); // hInternal - internals.push_back(shape(_batchSize, _hidSize)); // cInternal - internals.push_back(shape(_batchSize, 1)); // dummyOnes - internals.push_back(shape(_batchSize, 4*_hidSize)); // gates + // forward() allocates its own per-direction scratch, so no engine internals are needed + internals.clear(); return false; } @@ -201,7 +231,7 @@ class LSTM2LayerImpl CV_FINAL : public LSTM2Layer { CV_Assert(inputs[0] == CV_32F || inputs[0] == CV_64F); // Only floating-point types are supported currently outputs.assign(requiredOutputs, inputs[0]); - internals.assign(4, inputs[0]); + internals.clear(); } virtual int64 getFLOPS(const std::vector &inputs, @@ -223,19 +253,195 @@ class LSTM2LayerImpl CV_FINAL : public LSTM2Layer return (int64)numDirs * _seqLen * _batchSize * flopsPerStep; } + // Run one direction's recurrence into its own scratch so both directions run concurrently. + // Writes hOutAll columns [i*H, (i+1)*H) and the matching cOut slice. + void forwardDirection(int i, int numDirs, const Mat& xTs, const Mat& h0All, const Mat& c0All, + Mat& hOutAll, Mat& cOut) const + { + const int batchSizeTotal = seqLenth * batchSize; + const int dtype = xTs.type(); + + Mat Wx = weightBlobs[0].rowRange(i * weightBlobs[0].rows / numDirs, (i + 1) * weightBlobs[0].rows / numDirs); + Mat Wh = weightBlobs[1].rowRange(i * weightBlobs[1].rows / numDirs, (i + 1) * weightBlobs[1].rows / numDirs); + Mat bias = weightBlobs[2].colRange(i * weightBlobs[2].cols / numDirs, (i + 1) * weightBlobs[2].cols / numDirs); + Mat pI, pF, pO; + if (usePeephole) + { + pI = weightBlobs[3].rowRange(i * weightBlobs[3].rows / numDirs, (i + 1) * weightBlobs[3].rows / numDirs); + pF = weightBlobs[4].rowRange(i * weightBlobs[4].rows / numDirs, (i + 1) * weightBlobs[4].rows / numDirs); + pO = weightBlobs[5].rowRange(i * weightBlobs[5].rows / numDirs, (i + 1) * weightBlobs[5].rows / numDirs); + } + + Mat hInternal(batchSize, numHidden, dtype); + Mat cInternal(batchSize, numHidden, dtype); + h0All.rowRange(i * batchSize, (i + 1) * batchSize).copyTo(hInternal); + c0All.rowRange(i * batchSize, (i + 1) * batchSize).copyTo(cInternal); + + Mat hOutTs(batchSizeTotal, numHidden, dtype); + Mat cOutTs; + if (produceCellOutput) + { + cOutTs = cOut.reshape(1, batchSizeTotal); + cOutTs = cOutTs.colRange(i * cOutTs.cols / numDirs, (i + 1) * cOutTs.cols / numDirs); + } + + // Batched projection: gatesAll = Wx*x + bias for the whole sequence. + const int gateN = 4 * numHidden, projK = xTs.cols; + Mat gatesAll; + if (dtype == CV_32F && (int)packedWx.size() == numDirs && !packedWx[i].empty()) + { + FastGemmOpt opt = gemmOpt; + opt.multi_thread = (numDirs == 1); // directions are parallelized when bidirectional + repeat(bias, batchSizeTotal, 1, gatesAll); // each row starts as the bias + fastGemm(false, batchSizeTotal, gateN, projK, 1.f, xTs.ptr(), projK, + packedWx[i].data(), 1.f, gatesAll.ptr(), gateN, opt); + } + else // fp64 / unpacked fallback + { + Mat onesCol(batchSizeTotal, 1, dtype, Scalar(1)); + gatesAll.create(batchSizeTotal, gateN, dtype); + gemm(onesCol, bias, 1.0, noArray(), 0.0, gatesAll); + gemm(xTs, Wx, 1.0, gatesAll, 1.0, gatesAll, GEMM_2_T); + } + + // fastGEMM1T (recurrent step) needs contiguous fp32 operands +#if CV_TRY_AVX2 || CV_TRY_AVX + bool canUseAvxH = hInternal.isContinuous() && gatesAll.isContinuous() + && Wh.depth() == CV_32F && hInternal.depth() == CV_32F && gatesAll.depth() == CV_32F && Wh.cols >= 8; +#endif +#if CV_TRY_SVE + bool canUseSveH = hInternal.isContinuous() && gatesAll.isContinuous() + && Wh.depth() == CV_32F && hInternal.depth() == CV_32F && gatesAll.depth() == CV_32F; +#endif +#if CV_TRY_NEON + bool canUseNeonH = hInternal.isContinuous() && gatesAll.isContinuous() + && Wh.depth() == CV_32F && hInternal.depth() == CV_32F && gatesAll.depth() == CV_32F && Wh.cols >= 4; +#endif + + int tsStart, tsEnd, tsInc; + if (reverse || i == 1) { + tsStart = seqLenth - 1; + tsEnd = -1; + tsInc = -1; + } + else { + tsStart = 0; + tsEnd = seqLenth; + tsInc = 1; + } + + for (int ts = tsStart; ts != tsEnd; ts += tsInc) + { + Range curRowRange(ts*batchSize, (ts + 1)*batchSize); + Mat gates = gatesAll.rowRange(curRowRange); // already holds Wx * x_t + b + + // gates += Wh * h_{t-1} +#if CV_TRY_AVX2 + if (useAVX2 && canUseAvxH) + { + for (int n = 0; n < hInternal.rows; n++) + opt_AVX2::fastGEMM1T(hInternal.ptr(n), Wh.ptr(), Wh.step1(), + gates.ptr(n), gates.ptr(n), Wh.rows, Wh.cols); + } + else +#endif +#if CV_TRY_AVX + if (useAVX && canUseAvxH) + { + for (int n = 0; n < hInternal.rows; n++) + opt_AVX::fastGEMM1T(hInternal.ptr(n), Wh.ptr(), Wh.step1(), + gates.ptr(n), gates.ptr(n), Wh.rows, Wh.cols); + } + else +#endif +#if CV_TRY_SVE + if (useSVE && canUseSveH) + { + for (int n = 0; n < hInternal.rows; n++) + opt_SVE::fastGEMM1T(hInternal.ptr(n), Wh.ptr(), Wh.step1(), + gates.ptr(n), gates.ptr(n), Wh.rows, Wh.cols); + } + else +#endif +#if CV_TRY_NEON + if (useNEON && canUseNeonH) + { + for (int n = 0; n < hInternal.rows; n++) + opt_NEON::fastGEMM1T(hInternal.ptr(n), Wh.ptr(), Wh.step1(), + gates.ptr(n), gates.ptr(n), Wh.rows, Wh.cols); + } + else +#endif + { + gemm(hInternal, Wh, 1, gates, 1, gates, GEMM_2_T); + } + + Mat gateI = gates.colRange(0*numHidden, 1*numHidden); + Mat gateF = gates.colRange(1*numHidden, 2*numHidden); + Mat gateO = gates.colRange(2*numHidden, 3*numHidden); + Mat gateG = gates.colRange(3*numHidden, 4*numHidden); + + if (forgetBias){ + add(gateF, forgetBias, gateF); + } + + if (usePeephole) + { + Mat gatesIF = gates.colRange(0, 2*numHidden); + gemm(cInternal, pI, 1, gateI, 1, gateI); + gemm(cInternal, pF, 1, gateF, 1, gateF); + f_activation(gatesIF, gatesIF); + } + else + { + Mat gatesIFO = gates.colRange(0, 3*numHidden); + f_activation(gatesIFO, gatesIFO); + } + + g_activation(gateG, gateG); + + //compute c_t + multiply(gateF, cInternal, gateF); // f_t (*) c_{t-1} + multiply(gateI, gateG, gateI); // i_t (*) g_t + add(gateF, gateI, cInternal); // c_t = f_t (*) c_{t-1} + i_t (*) g_t + + if (useCellClip) + { + min(cInternal, cellClip, cInternal); + max(cInternal, -cellClip, cInternal); + } + + if (usePeephole) + { + gemm(cInternal, pO, 1, gateO, 1, gateO); + f_activation(gateO, gateO); + } + + //compute h_t + h_activation(cInternal, hInternal); + multiply(gateO, hInternal, hInternal); + + hInternal.copyTo(hOutTs.rowRange(curRowRange)); + + if (produceCellOutput) + cInternal.copyTo(cOutTs.rowRange(curRowRange)); + } + + // slice this direction's result into the assembly buffer + hOutTs.copyTo(hOutAll.colRange(i * numHidden, (i + 1) * numHidden)); + } + void forward(InputArrayOfArrays inputs_arr, OutputArrayOfArrays outputs_arr, OutputArrayOfArrays internals_arr) CV_OVERRIDE { - std::vector input, output, internals; + std::vector input, output; inputs_arr.getMatVector(input); outputs_arr.getMatVector(output); - internals_arr.getMatVector(internals); int numInputs = input.size(); int inpSize = input[0].size[2]; - int hidSize = numHidden; // determine seqLen and batchSize if (useTimestampDim) @@ -254,177 +460,54 @@ class LSTM2LayerImpl CV_FINAL : public LSTM2Layer batchSize = input[0].size[0]; } - // ONNX LSTM inputs: X(0), W(1), R(2), B(3), sequence_lens(4), - // initial_h(5), initial_c(6), P(7). Inputs 3..7 are optional and - // may be present-but-empty (e.g. CNTK exports keep all 8 slots with - // empties for unused ones). Gather by presence/non-emptiness rather - // than by the raw input count so optional/empty inputs are ignored. + // ONNX LSTM inputs 0..7: X,W,R,B,sequence_lens,initial_h,initial_c,P; slots 3..7 optional. + // Test non-emptiness, not count: some exports keep empty slots for unused inputs. auto hasInput = [&](int idx) { return idx < numInputs && !input[idx].empty(); }; - std::vector blobs_; - int hidShape [] = {1 + static_cast(bidirectional), batchSize, numHidden}; - int biasShape [] = {1 + static_cast(bidirectional), 8 * numHidden}; - CV_Assert(numInputs >= 3); // X, W, R are mandatory - blobs_.push_back(input[1].clone()); // W - blobs_.push_back(input[2].clone()); // R - // B - blobs_.push_back(hasInput(3) ? input[3] : Mat::zeros(2, biasShape, input[0].type())); - // initial_h - blobs_.push_back(hasInput(5) ? input[5] : Mat::zeros(3, hidShape, input[0].type())); - // initial_c - blobs_.push_back(hasInput(6) ? input[6] : Mat::zeros(3, hidShape, input[0].type())); - // P (peephole) - only when the layer was configured to use it and the input is present - if (usePeephole && hasInput(7)) - blobs_.push_back(input[7]); + + const int numDirs = 1 + static_cast(bidirectional); + const int batchSizeTotal = seqLenth * batchSize; + + // Weight transform is shape-independent: once for const weights, else every forward. + if (!constWeights || !weightsPacked) + { + packWeights(input); + weightsPacked = true; + } + + // Initial states are shape dependent, so (re)build them every forward. + Mat h0All = hasInput(5) ? input[5].reshape(1, input[5].size[0] * input[5].size[1]) + : Mat::zeros(numDirs * batchSize, numHidden, input[0].type()); + Mat c0All = hasInput(6) ? input[6].reshape(1, input[6].size[0] * input[6].size[1]) + : Mat::zeros(numDirs * batchSize, numHidden, input[0].type()); // set outputs to 0 for (auto& out : output) out.setTo(0); - - // convert weights to 2d matrices ease of use later in the forward pass - transformBlobs(blobs_); - - const int numDirs = 1 + static_cast(bidirectional); - const int batchSizeTotal = seqLenth * batchSize; - - Mat hInternal = internals[0], - cInternal = internals[1], - dummyOnes = internals[2], - gates = internals[3]; - - - Mat cOutTs; - // seq-major scratch for the cell states, (seq, batch, dirs, hid) like the recurrence writes + // seq-major cell-state scratch: (seq, batch, dirs, hid), matching the recurrence. int cOutShape[] = {seqLenth, batchSize, numDirs, numHidden}; Mat cOut = produceCellOutput ? Mat::zeros(4, cOutShape, output[0].type()) : Mat(); - Mat hOutTs = Mat::zeros(seqLenth * batchSize, hidSize, output[0].type()); Mat xTs = input[0].reshape(1, batchSizeTotal); - // seq-major assembly buffer for Y. The final result is transposed from it INTO output[0]: - // the preallocated output tensor must never be reallocated or get its header replaced, - // or the result would silently detach from the graph. + // seq-major Y assembly buffer; transposed into output[0] below. + // Never reallocate output[0]'s header or it detaches from the graph. Mat hOutAll(batchSizeTotal, numDirs * numHidden, output[0].type()); - // Initialize Wx, Wh, bias, h_0, c_0, pI, pF, pO - Mat Wx, Wh, bias, h_0, c_0, pI, pF, pO; - - for (int i = 0; i < numDirs; i++) + // Directions are independent (disjoint output columns): run them in parallel. + parallel_for_(Range(0, numDirs), [&](const Range& r) { - // slice required weights for each direction - Wx = blobs_[0].rowRange(i * blobs_[0].rows / numDirs, (i + 1) * blobs_[0].rows / numDirs); - Wh = blobs_[1].rowRange(i * blobs_[1].rows / numDirs, (i + 1) * blobs_[1].rows / numDirs); - bias = blobs_[2].colRange(i * blobs_[2].cols / numDirs, (i + 1) * blobs_[2].cols / numDirs); - h_0 = blobs_[3].rowRange(i * blobs_[3].rows / numDirs, (i + 1) * blobs_[3].rows / numDirs); - c_0 = blobs_[4].rowRange(i * blobs_[4].rows / numDirs, (i + 1) * blobs_[4].rows / numDirs); + for (int i = r.start; i < r.end; i++) + forwardDirection(i, numDirs, xTs, h0All, c0All, hOutAll, cOut); + }, numDirs); - if (usePeephole) - { - // slice required weights for each direction - pI = blobs_[5].rowRange(i * blobs_[5].rows / numDirs, (i + 1) * blobs_[5].rows / numDirs); - pF = blobs_[6].rowRange(i * blobs_[6].rows / numDirs, (i + 1) * blobs_[6].rows / numDirs); - pO = blobs_[7].rowRange(i * blobs_[7].rows / numDirs, (i + 1) * blobs_[7].rows / numDirs); - } - - h_0.copyTo(hInternal); - c_0.copyTo(cInternal); - dummyOnes.setTo(1.); - gates.setTo(0.); - - if (produceCellOutput) - { - cOutTs = cOut.reshape(1, batchSizeTotal); - cOutTs = cOutTs.colRange(i * cOutTs.cols / numDirs, (i + 1) * cOutTs.cols / numDirs); - } - - int tsStart, tsEnd, tsInc; - if (reverse || i == 1) { - tsStart = seqLenth - 1; - tsEnd = -1; - tsInc = -1; - } - else { - tsStart = 0; - tsEnd = seqLenth; - tsInc = 1; - } - - // main loop of LSTM forward pass - for (int ts = tsStart; ts != tsEnd; ts += tsInc) - { - Range curRowRange(ts*batchSize, (ts + 1)*batchSize); - Mat xCurr = xTs.rowRange(curRowRange); - - gemm(xCurr, Wx, 1, gates, 0, gates, GEMM_2_T); // Wx * x_t - gemm(dummyOnes, bias, 1, gates, 1, gates); //+b - gemm(hInternal, Wh, 1, gates, 1, gates, GEMM_2_T); //+Wh * h_{t-1} - - Mat gateI = gates.colRange(0*hidSize, 1*hidSize); - Mat gateF = gates.colRange(1*hidSize, 2*hidSize); - Mat gateO = gates.colRange(2*hidSize, 3*hidSize); - Mat gateG = gates.colRange(3*hidSize, 4*hidSize); - - if (forgetBias){ - add(gateF, forgetBias, gateF); - } - - if (usePeephole) - { - Mat gatesIF = gates.colRange(0, 2*hidSize); - gemm(cInternal, pI, 1, gateI, 1, gateI); - gemm(cInternal, pF, 1, gateF, 1, gateF); - f_activation(gatesIF, gatesIF); - } - else - { - Mat gatesIFO = gates.colRange(0, 3*hidSize); - f_activation(gatesIFO, gatesIFO); - } - - g_activation(gateG, gateG); - - //compute c_t - multiply(gateF, cInternal, gateF); // f_t (*) c_{t-1} - multiply(gateI, gateG, gateI); // i_t (*) g_t - add(gateF, gateI, cInternal); // c_t = f_t (*) c_{t-1} + i_t (*) g_t - - if (useCellClip) - { - min(cInternal, cellClip, cInternal); - max(cInternal, -cellClip, cInternal); - } - - if (usePeephole) - { - gemm(cInternal, pO, 1, gateO, 1, gateO); - f_activation(gateO, gateO); - } - - //compute h_t - h_activation(cInternal, hInternal); - multiply(gateO, hInternal, hInternal); - - //save results in output blobs - hInternal.copyTo(hOutTs.rowRange(curRowRange)); - - if (produceCellOutput) - cInternal.copyTo(cOutTs.rowRange(curRowRange)); - - } - - // slice in the result from each direction to the assembly buffer - hOutTs.copyTo(hOutAll.colRange(i * hOutTs.cols, (i + 1) * hOutTs.cols)); - } - - // (seq*batch, dirs*hid) -> (seq, batch, dirs, hid), then into the ONNX Y layout for this - // `layout` attribute, written INTO the preallocated output[0] (transposeND's exact-shape - // create() keeps it in place) + // Reshape to (seq, batch, dirs, hid), then transpose into output[0] per `layout`. int shp1[] = {seqLenth, batchSize, numDirs, numHidden}; Mat y4d = hOutAll.reshape(1, sizeof(shp1)/sizeof(shp1[0]), shp1); - Mat ySeqFirst; // (seq, dirs, batch, hid) - the layout=0 Y; Yh is sliced from it + Mat ySeqFirst; // (seq, dirs, batch, hid): the layout=0 Y; Yh sliced from it if (layout == SEQ_BATCH_HID) { cv::transposeND(y4d, {0, 2, 1, 3}, output[0]); ySeqFirst = output[0]; @@ -541,78 +624,68 @@ class LSTM2LayerImpl CV_FINAL : public LSTM2Layer } } - void transformBlobs(std::vector& blobs) + // Fill weightBlobs: [0]=Wx, [1]=Wh, [2]=bias, and (peephole) [3]=pI, [4]=pF, [5]=pO. + // Shape-independent, so it runs once for constant weights. + void packWeights(const std::vector& input) { - Mat &Wx = blobs[0]; - Mat &Wh = blobs[1]; - Mat &b = blobs[2]; + int biasShape[] = {1 + static_cast(bidirectional), 8 * numHidden}; - const int numHidden = Wh.size[2]; + weightBlobs.clear(); + weightBlobs.push_back(input[1].clone()); // W -> Wx + weightBlobs.push_back(input[2].clone()); // R -> Wh + weightBlobs.push_back((input.size() > 3 && !input[3].empty()) + ? input[3].clone() + : Mat::zeros(2, biasShape, input[0].type())); // B + bool hasP = usePeephole && input.size() > 7 && !input[7].empty(); + if (hasP) + weightBlobs.push_back(input[7].clone()); // P - Mat h0, c0; - // check weather input is dynamic or not: hx, cx are given by user. - // Resahpe if only they are given - if (!blobs[3].empty()){ - h0 = blobs[3]; - h0 = h0.reshape(1, h0.size[0] * h0.size[1]); - } - if (!blobs[4].empty()){ - c0 = blobs[4]; - c0 = c0.reshape(1, c0.size[0] * c0.size[1]); - } + Mat& Wx = weightBlobs[0]; + Mat& Wh = weightBlobs[1]; + Mat& b = weightBlobs[2]; b = b.reshape(1, b.size[0]); Mat bx = b.colRange(0, b.cols / 2); Mat bh = b.colRange(b.cols / 2, b.cols); b = bx + bh; - auto toIFOC = [] (Mat& in) { - int first = in.size[0]; - int rest = in.total() / first / 4; - // every weight blob contains weights for Input, Output, Forget and Cell gates - Mat m = in.reshape(1, {first, 4, rest}); - Mat outputGate = m.col(1); - Mat forgetGate = m.col(2); - std::swap_ranges(outputGate.begin(), outputGate.end(), forgetGate.begin()); - }; + reorderGatesIOFCtoIFOC(Wx); + reorderGatesIOFCtoIFOC(Wh); + reorderGatesIOFCtoIFOC(b); - toIFOC(Wx); - toIFOC(Wh); - toIFOC(b); + weightBlobs[0] = Wx.reshape(1, Wx.size[0] * Wx.size[1]); + weightBlobs[1] = Wh.reshape(1, Wh.size[0] * Wh.size[1]); + weightBlobs[2] = b.reshape(1, 1); - - Wx = Wx.reshape(1, Wx.size[0] * Wx.size[1]); - Wh = Wh.reshape(1, Wh.size[0] * Wh.size[1]); - - - blobs[0] = Wx; - blobs[1] = Wh; - blobs[2] = b.reshape(1, 1); - - if (!blobs[3].empty()){ - blobs[3] = h0; - } - if (!blobs[4].empty()){ - blobs[4] = c0; + // Prepack Wx per direction (fp32) so the projection is one batched GEMM. + packedWx.clear(); + if (weightBlobs[0].type() == CV_32F) + { + const int numDirs = 1 + static_cast(bidirectional); + const int N = weightBlobs[0].rows / numDirs; // 4 * numHidden + packedWx.resize(numDirs); + for (int d = 0; d < numDirs; d++) + { + Mat WxDir = weightBlobs[0].rowRange(d * N, (d + 1) * N); + fastGemmPackB(WxDir, packedWx[d], true, gemmOpt); + } } - if (blobs.size() == 5) { + if (!hasP) return; - } - Mat P = blobs[5]; - blobs[5] = P.colRange(0, numHidden); - blobs[5] = blobs[5].clone().reshape(1, blobs[5].total()); // Single column. - blobs[5] = Mat::diag(blobs[5]); + Mat P = weightBlobs[3]; + weightBlobs[3] = P.colRange(0, numHidden); + weightBlobs[3] = weightBlobs[3].clone().reshape(1, weightBlobs[3].total()); // Single column. + weightBlobs[3] = Mat::diag(weightBlobs[3]); - blobs.push_back(P.colRange(numHidden, 2 * numHidden)); - blobs[6] = blobs[6].clone().reshape(1, blobs[6].total()); // Single column. - blobs[6] = Mat::diag(blobs[6]); + weightBlobs.push_back(P.colRange(numHidden, 2 * numHidden)); + weightBlobs[4] = weightBlobs[4].clone().reshape(1, weightBlobs[4].total()); // Single column. + weightBlobs[4] = Mat::diag(weightBlobs[4]); - blobs.push_back(P.colRange(2 * numHidden, 3 * numHidden)); - blobs[7] = blobs[7].clone().reshape(1, blobs[7].total()); // Single column. - blobs[7] = Mat::diag(blobs[7]); - return; + weightBlobs.push_back(P.colRange(2 * numHidden, 3 * numHidden)); + weightBlobs[5] = weightBlobs[5].clone().reshape(1, weightBlobs[5].total()); // Single column. + weightBlobs[5] = Mat::diag(weightBlobs[5]); } }; diff --git a/modules/dnn/src/onnx/onnx_importer2.cpp b/modules/dnn/src/onnx/onnx_importer2.cpp index 3b63e59114..4fc7ba1884 100644 --- a/modules/dnn/src/onnx/onnx_importer2.cpp +++ b/modules/dnn/src/onnx/onnx_importer2.cpp @@ -1353,6 +1353,14 @@ void ONNXImporter2::parseLSTM(LayerParams& layerParams, const opencv_onnx::NodeP if (lstm_proto.input_size() == 8 && !lstm_proto.input(7).empty()) layerParams.set("use_peephole", true); + // W/R (and B/P if present) are usually const: let the layer transform them once. + bool const_weights = lstm_proto.input_size() >= 3 && + net.isConstArg(node_inputs[1]) && net.isConstArg(node_inputs[2]); + if (const_weights && lstm_proto.input_size() > 3 && !lstm_proto.input(3).empty()) + const_weights = net.isConstArg(node_inputs[3]); + if (const_weights && lstm_proto.input_size() > 7 && !lstm_proto.input(7).empty()) + const_weights = net.isConstArg(node_inputs[7]); + layerParams.set("const_weights", const_weights); addLayer(layerParams, lstm_proto); }