| 16 | public: |
| 17 | virtual ~UnstackTest() = default; |
| 18 | virtual bool run(int precision) { |
| 19 | auto input = _Input({3, 1, 2, 2}, NCHW); |
| 20 | input->setName("input_tensor"); |
| 21 | // set input data |
| 22 | const float inpudata[] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0}; |
| 23 | auto inputPtr = input->writeMap<float>(); |
| 24 | memcpy(inputPtr, inpudata, 12 * sizeof(float)); |
| 25 | input->unMap(); |
| 26 | auto outputs = _Unstack(input, -2); |
| 27 | const std::vector<float> expectedOutput_0 = {1.0, 2.0, 5.0, 6.0, 9.0, 10.0}; |
| 28 | const std::vector<float> expectedOutput_1 = {3.0, 4.0, 7.0, 8.0, 11.0, 12.0}; |
| 29 | auto gotOutput_0 = outputs[0]->readMap<float>(); |
| 30 | auto gotOutput_1 = outputs[1]->readMap<float>(); |
| 31 | if (!checkVector<float>(gotOutput_0, expectedOutput_0.data(), 6, 0.01)) { |
| 32 | MNN_ERROR("UnstackTest test failed!\n"); |
| 33 | return false; |
| 34 | } |
| 35 | if (!checkVector<float>(gotOutput_1, expectedOutput_1.data(), 6, 0.01)) { |
| 36 | MNN_ERROR("UnstackTest test failed!\n"); |
| 37 | return false; |
| 38 | } |
| 39 | return true; |
| 40 | } |
| 41 | }; |
| 42 | MNNTestSuiteRegister(UnstackTest, "op/unstack"); |