1d support for einsum

pull/24830/head
Abduragim 1 year ago
parent bae435a5a7
commit 6c28d7140a
  1. 4
      modules/dnn/src/layers/einsum_layer.cpp
  2. 6
      modules/dnn/test/test_onnx_importer.cpp

@ -105,7 +105,7 @@ static Mat batchwiseMatMul(
// input1 should of size MxK // input1 should of size MxK
// check if input1 needs reshape, if need reshape // check if input1 needs reshape, if need reshape
if (input1.dims > 2 || input1.size[0] != M || input1.size[1] != K) if (input1.dims > 2 || input1.size[0] != M || (input1.dims > 1 && input1.size[1] != K) || input1.dims == 1)
{ {
int shape[] = {static_cast<int>(M), static_cast<int>(K)}; int shape[] = {static_cast<int>(M), static_cast<int>(K)};
reshapedInput1 = input1.reshape(1, 2, shape); reshapedInput1 = input1.reshape(1, 2, shape);
@ -113,7 +113,7 @@ static Mat batchwiseMatMul(
// input2 should be of size KxN // input2 should be of size KxN
// check if input2 needs reshape, if needs reshape // check if input2 needs reshape, if needs reshape
if (input2.dims > 2 || input2.size[0] != K || input2.size[1] != N) if (input2.dims > 2 || input2.size[0] != K || (input2.dims > 1 && input2.size[1] != N) || input2.dims == 1)
{ {
int shape2[] = {static_cast<int>(K), static_cast<int>(N)}; int shape2[] = {static_cast<int>(K), static_cast<int>(N)};
reshapedInput2 = input2.reshape(1, 2, shape2); reshapedInput2 = input2.reshape(1, 2, shape2);

@ -1452,7 +1452,7 @@ TEST_P(Test_ONNX_layers, LSTM_layout_batch)
testONNXModels("lstm_layout_1", npy, 0.005, 0.005, false, false, 3); testONNXModels("lstm_layout_1", npy, 0.005, 0.005, false, false, 3);
} }
TEST_P(Test_ONNX_layers, DISABLED_Einsum_1D) TEST_P(Test_ONNX_layers, Einsum_1D)
{ {
testONNXModels("einsum_1d", npy, 0, 0, false, false, 2); testONNXModels("einsum_1d", npy, 0, 0, false, false, 2);
} }
@ -1482,12 +1482,12 @@ TEST_P(Test_ONNX_layers, Einsum_5D)
testONNXModels("einsum_5d", npy, 0, 0, false, false, 2); testONNXModels("einsum_5d", npy, 0, 0, false, false, 2);
} }
TEST_P(Test_ONNX_layers, DISABLED_Einsum_InnerProduct) TEST_P(Test_ONNX_layers, Einsum_InnerProduct)
{ {
testONNXModels("einsum_inner", npy, 0, 0, false, false, 2); testONNXModels("einsum_inner", npy, 0, 0, false, false, 2);
} }
TEST_P(Test_ONNX_layers, DISABLED_Einsum_HadamardProduct) TEST_P(Test_ONNX_layers, Einsum_HadamardProduct)
{ {
testONNXModels("einsum_hadamard", npy, 0, 0, false, false, 2); testONNXModels("einsum_hadamard", npy, 0, 0, false, false, 2);
} }

Loading…
Cancel
Save