From 4dbfd7193ec3fdcce1017225517dea4ee9c9b702 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E4=BD=B0=E9=98=85?= <43716063+Baiyuetribe@users.noreply.github.com> Date: Thu, 16 Jan 2025 21:10:51 +0800 Subject: [PATCH] check c++03 and ctest --- .github/workflows/linux-x64-cpu-gcc.yml | 218 +++++++------- src/CMakeLists.txt | 1 + src/layer/indexselect.cpp | 270 ++++++++++++++++++ src/layer/indexselect.h | 37 +++ tests/CMakeLists.txt | 1 + tests/test_indexselect.cpp | 113 ++++++++ tools/pnnx/tests/ncnn/CMakeLists.txt | 1 + .../tests/ncnn/test_torch_index_select.py | 149 ++++++++++ 8 files changed, 681 insertions(+), 109 deletions(-) create mode 100644 src/layer/indexselect.cpp create mode 100644 src/layer/indexselect.h create mode 100644 tests/test_indexselect.cpp create mode 100644 tools/pnnx/tests/ncnn/test_torch_index_select.py diff --git a/.github/workflows/linux-x64-cpu-gcc.yml b/.github/workflows/linux-x64-cpu-gcc.yml index ab2185be3e7..31abbe47c25 100644 --- a/.github/workflows/linux-x64-cpu-gcc.yml +++ b/.github/workflows/linux-x64-cpu-gcc.yml @@ -1,33 +1,33 @@ name: linux-x64-cpu-gcc on: push: - branches: [master] + # branches: [master] paths: - - '.github/workflows/linux-x64-cpu-gcc.yml' - - 'toolchains/host-c.gcc.toolchain.cmake' - - 'CMakeLists.txt' - - 'cmake/**' - - 'src/*' - - 'src/layer/*' - - 'src/layer/x86/**' - - 'tests/**' - - 'tools/**' - - '!tools/pnnx/**' - - 'examples/**' + - ".github/workflows/linux-x64-cpu-gcc.yml" + - "toolchains/host-c.gcc.toolchain.cmake" + - "CMakeLists.txt" + - "cmake/**" + - "src/*" + - "src/layer/*" + - "src/layer/x86/**" + - "tests/**" + - "tools/**" + - "!tools/pnnx/**" + - "examples/**" pull_request: branches: [master] paths: - - '.github/workflows/linux-x64-cpu-gcc.yml' - - 'toolchains/host-c.gcc.toolchain.cmake' - - 'CMakeLists.txt' - - 'cmake/**' - - 'src/*' - - 'src/layer/*' - - 'src/layer/x86/**' - - 'tests/**' - - 'tools/**' - - '!tools/pnnx/**' - - 'examples/**' + - ".github/workflows/linux-x64-cpu-gcc.yml" + - "toolchains/host-c.gcc.toolchain.cmake" + - "CMakeLists.txt" + - "cmake/**" + - "src/*" + - "src/layer/*" + - "src/layer/x86/**" + - "tests/**" + - "tools/**" + - "!tools/pnnx/**" + - "examples/**" concurrency: group: linux-x64-cpu-gcc-${{ github.ref }} cancel-in-progress: true @@ -38,97 +38,97 @@ jobs: linux-gcc: runs-on: ubuntu-20.04 steps: - - uses: actions/checkout@v4 - - name: update - run: sudo apt-get update - - name: protobuf - run: sudo apt-get install libprotobuf-dev protobuf-compiler libopencv-dev - - name: build-sse2 - run: | - mkdir build-sse2 && cd build-sse2 - cmake -DNCNN_AVX=OFF -DNCNN_AVX2=OFF -DNCNN_BUILD_TESTS=ON .. - cmake --build . -j $(nproc) - - name: test-sse2 - run: cd build-sse2 && ctest --output-on-failure -j $(nproc) - - name: build-shared - run: | - mkdir build-shared && cd build-shared - cmake -DNCNN_AVX2=ON -DNCNN_SHARED_LIB=ON .. - cmake --build . -j $(nproc) - - name: build-avx2 - run: | - mkdir build-avx2 && cd build-avx2 - cmake -DNCNN_AVX2=ON -DNCNN_BUILD_TESTS=ON .. - cmake --build . -j $(nproc) - - name: test-avx2 - run: cd build-avx2 && ctest --output-on-failure -j $(nproc) - - name: build-avx - run: | - mkdir build-avx && cd build-avx - cmake -DNCNN_AVX2=OFF -DNCNN_AVX=ON -DNCNN_BUILD_TESTS=ON .. - cmake --build . -j $(nproc) - - name: test-avx - run: cd build-avx && ctest --output-on-failure -j $(nproc) - - name: build-avx1-2 - run: | - mkdir build-avx1-2 && cd build-avx1-2 - cmake -DNCNN_AVX2=ON -DNCNN_AVX=ON -DNCNN_BUILD_TESTS=ON .. - cmake --build . -j $(nproc) - - name: test-avx1-2 - run: cd build-avx1-2 && ctest --output-on-failure -j $(nproc) - - name: build-noint8 - run: | - mkdir build-noint8 && cd build-noint8 - cmake -DNCNN_INT8=OFF -DNCNN_BUILD_TESTS=ON .. - cmake --build . -j $(nproc) - - name: test-noint8 - run: cd build-noint8 && ctest --output-on-failure -j $(nproc) + - uses: actions/checkout@v4 + - name: update + run: sudo apt-get update + - name: protobuf + run: sudo apt-get install libprotobuf-dev protobuf-compiler libopencv-dev + - name: build-sse2 + run: | + mkdir build-sse2 && cd build-sse2 + cmake -DNCNN_AVX=OFF -DNCNN_AVX2=OFF -DNCNN_BUILD_TESTS=ON .. + cmake --build . -j $(nproc) + - name: test-sse2 + run: cd build-sse2 && ctest --output-on-failure -j $(nproc) + - name: build-shared + run: | + mkdir build-shared && cd build-shared + cmake -DNCNN_AVX2=ON -DNCNN_SHARED_LIB=ON .. + cmake --build . -j $(nproc) + - name: build-avx2 + run: | + mkdir build-avx2 && cd build-avx2 + cmake -DNCNN_AVX2=ON -DNCNN_BUILD_TESTS=ON .. + cmake --build . -j $(nproc) + - name: test-avx2 + run: cd build-avx2 && ctest --output-on-failure -j $(nproc) + - name: build-avx + run: | + mkdir build-avx && cd build-avx + cmake -DNCNN_AVX2=OFF -DNCNN_AVX=ON -DNCNN_BUILD_TESTS=ON .. + cmake --build . -j $(nproc) + - name: test-avx + run: cd build-avx && ctest --output-on-failure -j $(nproc) + - name: build-avx1-2 + run: | + mkdir build-avx1-2 && cd build-avx1-2 + cmake -DNCNN_AVX2=ON -DNCNN_AVX=ON -DNCNN_BUILD_TESTS=ON .. + cmake --build . -j $(nproc) + - name: test-avx1-2 + run: cd build-avx1-2 && ctest --output-on-failure -j $(nproc) + - name: build-noint8 + run: | + mkdir build-noint8 && cd build-noint8 + cmake -DNCNN_INT8=OFF -DNCNN_BUILD_TESTS=ON .. + cmake --build . -j $(nproc) + - name: test-noint8 + run: cd build-noint8 && ctest --output-on-failure -j $(nproc) linux-gcc-cpp03-nostdio-nostring-simplestl: runs-on: ubuntu-20.04 steps: - - uses: actions/checkout@v4 - - name: build-nostdio - run: | - mkdir build-nostdio && cd build-nostdio - cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/host.gcc-c++03.toolchain.cmake -DNCNN_BUILD_TESTS=ON -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF .. - cmake --build . -j $(nproc) - - name: test-nostdio - run: cd build-nostdio && ctest --output-on-failure -j $(nproc) - - name: build-nostdio-nostring - run: | - mkdir build-nostdio-nostring && cd build-nostdio-nostring - cmake -DNCNN_STDIO=OFF -DNCNN_STRING=OFF -DNCNN_BUILD_TESTS=OFF -DNCNN_BUILD_BENCHMARK=OFF -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF .. - cmake --build . -j $(nproc) - - name: build-simplestl - run: | - mkdir build-simplestl && cd build-simplestl - cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/host-c.gcc.toolchain.cmake -DNCNN_STDIO=ON -DNCNN_STRING=ON -DNCNN_SIMPLESTL=ON -DNCNN_BUILD_TESTS=ON -DNCNN_BUILD_BENCHMARK=OFF -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF .. - cmake --build . -j $(nproc) - - name: test-simplestl - run: cd build-simplestl && ctest --output-on-failure -j $(nproc) - - name: build-simplestl-simpleomp - run: | - mkdir build-simplestl-simpleomp && cd build-simplestl-simpleomp - cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/host-c.gcc.toolchain.cmake -DNCNN_STDIO=ON -DNCNN_STRING=ON -DNCNN_SIMPLESTL=ON -DNCNN_SIMPLEOMP=ON -DNCNN_BUILD_TESTS=ON -DNCNN_BUILD_BENCHMARK=OFF -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF .. - cmake --build . -j $(nproc) - - name: test-simplestl-simpleomp - run: cd build-simplestl-simpleomp && ctest --output-on-failure -j $(nproc) + - uses: actions/checkout@v4 + - name: build-nostdio + run: | + mkdir build-nostdio && cd build-nostdio + cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/host.gcc-c++03.toolchain.cmake -DNCNN_BUILD_TESTS=ON -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF .. + cmake --build . -j $(nproc) + - name: test-nostdio + run: cd build-nostdio && ctest --output-on-failure -j $(nproc) + - name: build-nostdio-nostring + run: | + mkdir build-nostdio-nostring && cd build-nostdio-nostring + cmake -DNCNN_STDIO=OFF -DNCNN_STRING=OFF -DNCNN_BUILD_TESTS=OFF -DNCNN_BUILD_BENCHMARK=OFF -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF .. + cmake --build . -j $(nproc) + - name: build-simplestl + run: | + mkdir build-simplestl && cd build-simplestl + cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/host-c.gcc.toolchain.cmake -DNCNN_STDIO=ON -DNCNN_STRING=ON -DNCNN_SIMPLESTL=ON -DNCNN_BUILD_TESTS=ON -DNCNN_BUILD_BENCHMARK=OFF -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF .. + cmake --build . -j $(nproc) + - name: test-simplestl + run: cd build-simplestl && ctest --output-on-failure -j $(nproc) + - name: build-simplestl-simpleomp + run: | + mkdir build-simplestl-simpleomp && cd build-simplestl-simpleomp + cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/host-c.gcc.toolchain.cmake -DNCNN_STDIO=ON -DNCNN_STRING=ON -DNCNN_SIMPLESTL=ON -DNCNN_SIMPLEOMP=ON -DNCNN_BUILD_TESTS=ON -DNCNN_BUILD_BENCHMARK=OFF -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF .. + cmake --build . -j $(nproc) + - name: test-simplestl-simpleomp + run: cd build-simplestl-simpleomp && ctest --output-on-failure -j $(nproc) linux-gcc-avx512: runs-on: [self-hosted, linux, t4] steps: - - uses: actions/checkout@v4 - - name: build - env: - CC: gcc - CXX: g++ - LD_LIBRARY_PATH: /data/action/install/lib64 - run: | - mkdir build && cd build - cmake -DNCNN_AVX2=ON -DNCNN_AVX512=ON -DNCNN_AVX512VNNI=ON -DNCNN_BUILD_TESTS=ON -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF .. - cmake --build . -j 4 - - name: test - env: - LD_LIBRARY_PATH: /data/action/install/lib64 - run: cd build && ctest --output-on-failure -j 4 + - uses: actions/checkout@v4 + - name: build + env: + CC: gcc + CXX: g++ + LD_LIBRARY_PATH: /data/action/install/lib64 + run: | + mkdir build && cd build + cmake -DNCNN_AVX2=ON -DNCNN_AVX512=ON -DNCNN_AVX512VNNI=ON -DNCNN_BUILD_TESTS=ON -DNCNN_BUILD_TOOLS=OFF -DNCNN_BUILD_EXAMPLES=OFF .. + cmake --build . -j 4 + - name: test + env: + LD_LIBRARY_PATH: /data/action/install/lib64 + run: cd build && ctest --output-on-failure -j 4 diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c97235d97a0..33f9db2a630 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -169,6 +169,7 @@ ncnn_add_layer(Shrink) ncnn_add_layer(RMSNorm) ncnn_add_layer(Spectrogram) ncnn_add_layer(InverseSpectrogram) +ncnn_add_layer(IndexSelect) if(NCNN_VULKAN) ncnn_add_shader(${CMAKE_CURRENT_SOURCE_DIR}/convert_ycbcr.comp) diff --git a/src/layer/indexselect.cpp b/src/layer/indexselect.cpp new file mode 100644 index 00000000000..45dfc63b791 --- /dev/null +++ b/src/layer/indexselect.cpp @@ -0,0 +1,270 @@ +// Tencent is pleased to support the open source community by making ncnn available. +// +// Copyright (C) 2017 THL A29 Limited, a Tencent company. All rights reserved. +// +// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except +// in compliance with the License. You may obtain a copy of the License at +// +// https://opensource.org/licenses/BSD-3-Clause +// +// Unless required by applicable law or agreed to in writing, software distributed +// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR +// CONDITIONS OF ANY KIND, either express or implied. See the License for the +// specific language governing permissions and limitations under the License. + +#include "indexselect.h" + +namespace ncnn { +IndexSelect::IndexSelect() +{ + one_blob_only = false; // 是否单一输入 + support_inplace = false; // 是否支持原地运算 +} + +int IndexSelect::load_param(const ParamDict& pd) +{ + dim = pd.get(0, -1); // dim = [-dim~dim-1] + return 0; +} + +int IndexSelect::forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const +{ + const Mat& bottom_blob = bottom_blobs[0]; + const Mat& index_blob = bottom_blobs[1]; + Mat& top_blob = top_blobs[0]; // 仅1个输出 + int dims = bottom_blob.dims; + int w = bottom_blob.w; + int h = bottom_blob.h; + int d = bottom_blob.d; + int channels = bottom_blob.c; + size_t elemsize = bottom_blob.elemsize; + + int index_len = index_blob.w; // 索引数据 + + int axis = dim < 0 ? dim + dims : dim; + // 检查k值是否有效 + if (index_len < 1 || axis >= dims) + { + return -1; + } + + if (dims == 1) + { + // 创建输出blob + top_blob.create(index_len, elemsize, opt.blob_allocator); + const float* ptr = bottom_blob; + float* outptr = top_blob; + const int* index_ptr = index_blob; + for (int i = 0; i < index_len; i++) + { + outptr[i] = ptr[index_ptr[i]]; + } + } + else if (dims == 2) + { + if (axis == 0) + { + top_blob.create(w, index_len, elemsize, opt.blob_allocator); + const int* index_ptr = index_blob; + + for (int i = 0; i < index_len; i++) + { + int index = index_ptr[i]; + const float* ptr_row = bottom_blob.row(index); + float* outptr_row = top_blob.row(i); + memcpy(outptr_row, ptr_row, w * sizeof(float)); + } + } + else if (axis == 1) + { + top_blob.create(index_len, h, elemsize, opt.blob_allocator); + const int* index_ptr = index_blob; + for (int i = 0; i < h; i++) + { + const float* ptr_row = bottom_blob.row(i); + float* outptr_row = top_blob.row(i); + + // 对每一行,根据索引选择对应列 + for (int j = 0; j < index_len; j++) + { + int index = index_ptr[j]; + outptr_row[j] = ptr_row[index]; + } + } + } + } + else if (dims == 3) + { + if (axis == 0) // channels维度 + { + top_blob.create(w, h, index_len, elemsize, opt.blob_allocator); + const int* index_ptr = index_blob; + + for (int q = 0; q < index_len; q++) + { + int index = index_ptr[q]; + const Mat bottom_channel = bottom_blob.channel(index); + Mat top_channel = top_blob.channel(q); + + for (int i = 0; i < h; i++) + { + const float* ptr = bottom_channel.row(i); + float* outptr = top_channel.row(i); + memcpy(outptr, ptr, w * sizeof(float)); + } + } + } + else if (axis == 1) // h维度 + { + top_blob.create(w, index_len, channels, elemsize, opt.blob_allocator); + const int* index_ptr = index_blob; + + for (int q = 0; q < channels; q++) + { + const Mat bottom_channel = bottom_blob.channel(q); + Mat top_channel = top_blob.channel(q); + + for (int i = 0; i < index_len; i++) + { + int index = index_ptr[i]; + const float* ptr = bottom_channel.row(index); + float* outptr = top_channel.row(i); + memcpy(outptr, ptr, w * sizeof(float)); + } + } + } + else if (axis == 2) // w维度 + { + top_blob.create(index_len, h, channels, elemsize, opt.blob_allocator); + const int* index_ptr = index_blob; + + for (int q = 0; q < channels; q++) + { + const Mat bottom_channel = bottom_blob.channel(q); + Mat top_channel = top_blob.channel(q); + + for (int i = 0; i < h; i++) + { + const float* ptr = bottom_channel.row(i); + float* outptr = top_channel.row(i); + for (int j = 0; j < index_len; j++) + { + int index = index_ptr[j]; + outptr[j] = ptr[index]; + } + } + } + } + } + + else if (dims == 4) + { + if (axis == 0) // channels维度 + { + top_blob.create(w, h, d, index_len, elemsize, opt.blob_allocator); + const int* index_ptr = index_blob; + + for (int q = 0; q < index_len; q++) + { + int index = index_ptr[q]; + const Mat bottom_c = bottom_blob.channel(index); + Mat top_c = top_blob.channel(q); + + for (int z = 0; z < d; z++) + { + const Mat bottom_d = bottom_c.channel(z); + Mat top_d = top_c.channel(z); + for (int i = 0; i < h; i++) + { + const float* ptr = bottom_d.row(i); + float* outptr = top_d.row(i); + memcpy(outptr, ptr, w * sizeof(float)); + } + } + } + } + else if (axis == 1) // d维度 + { + top_blob.create(w, h, index_len, channels, elemsize, opt.blob_allocator); + const int* index_ptr = index_blob; + + for (int q = 0; q < channels; q++) + { + const Mat bottom_c = bottom_blob.channel(q); + Mat top_c = top_blob.channel(q); + + for (int z = 0; z < index_len; z++) + { + int index = index_ptr[z]; + const Mat bottom_d = bottom_c.channel(index); + Mat top_d = top_c.channel(z); + for (int i = 0; i < h; i++) + { + const float* ptr = bottom_d.row(i); + float* outptr = top_d.row(i); + memcpy(outptr, ptr, w * sizeof(float)); + } + } + } + } + else if (axis == 2) // h维度 + { + top_blob.create(w, index_len, d, channels, elemsize, opt.blob_allocator); + const int* index_ptr = index_blob; + + for (int q = 0; q < channels; q++) + { + const Mat bottom_c = bottom_blob.channel(q); + Mat top_c = top_blob.channel(q); + + for (int z = 0; z < d; z++) + { + const Mat bottom_d = bottom_c.channel(z); + Mat top_d = top_c.channel(z); + for (int i = 0; i < index_len; i++) + { + int index = index_ptr[i]; + const float* ptr = bottom_d.row(index); + float* outptr = top_d.row(i); + memcpy(outptr, ptr, w * sizeof(float)); + } + } + } + } + else if (axis == 3) // w维度 + { + top_blob.create(index_len, h, d, channels, elemsize, opt.blob_allocator); + const int* index_ptr = index_blob; + + for (int q = 0; q < channels; q++) + { + const Mat bottom_c = bottom_blob.channel(q); + Mat top_c = top_blob.channel(q); + + for (int z = 0; z < d; z++) + { + const Mat bottom_d = bottom_c.channel(z); + Mat top_d = top_c.channel(z); + for (int i = 0; i < h; i++) + { + const float* ptr = bottom_d.row(i); + float* outptr = top_d.row(i); + for (int j = 0; j < index_len; j++) + { + int index = index_ptr[j]; + outptr[j] = ptr[index]; + } + } + } + } + } + } + else + { + return -1; + } + + return 0; +} + +} // namespace ncnn \ No newline at end of file diff --git a/src/layer/indexselect.h b/src/layer/indexselect.h new file mode 100644 index 00000000000..1a0da571208 --- /dev/null +++ b/src/layer/indexselect.h @@ -0,0 +1,37 @@ +// Tencent is pleased to support the open source community by making ncnn available. +// +// Copyright (C) 2017 THL A29 Limited, a Tencent company. All rights reserved. +// +// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except +// in compliance with the License. You may obtain a copy of the License at +// +// https://opensource.org/licenses/BSD-3-Clause +// +// Unless required by applicable law or agreed to in writing, software distributed +// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR +// CONDITIONS OF ANY KIND, either express or implied. See the License for the +// specific language governing permissions and limitations under the License. + +#ifndef LAYER_INDEXSELECT_H +#define LAYER_INDEXSELECT_H + +#include "layer.h" + +namespace ncnn { + +class IndexSelect : public Layer +{ +public: + IndexSelect(); + + virtual int load_param(const ParamDict& pd); + + virtual int forward(const std::vector& bottom_blobs, std::vector& top_blobs, const Option& opt) const; + +public: + int dim; +}; + +} // namespace ncnn + +#endif // LAYER_TOPK_H diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index f55859e736e..4f44e8c9825 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -115,6 +115,7 @@ ncnn_add_layer_test(GRU) ncnn_add_layer_test(HardSigmoid) ncnn_add_layer_test(HardSwish) ncnn_add_layer_test(InnerProduct) +ncnn_add_layer_test(IndexSelect) ncnn_add_layer_test(InstanceNorm) ncnn_add_layer_test(Interp) ncnn_add_layer_test(InverseSpectrogram) diff --git a/tests/test_indexselect.cpp b/tests/test_indexselect.cpp new file mode 100644 index 00000000000..9577df81346 --- /dev/null +++ b/tests/test_indexselect.cpp @@ -0,0 +1,113 @@ +// Tencent is pleased to support the open source community by making ncnn available. +// +// Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved. +// +// Licensed under the BSD 3-Clause License (the "License"); you may not use this file except +// in compliance with the License. You may obtain a copy of the License at +// +// https://opensource.org/licenses/BSD-3-Clause +// +// Unless required by applicable law or agreed to in writing, software distributed +// under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR +// CONDITIONS OF ANY KIND, either express or implied. See the License for the +// specific language governing permissions and limitations under the License. + +#include "testutil.h" + +static ncnn::Mat IntArrayMat(int a0) +{ + ncnn::Mat m(1); + int* p = m; + p[0] = a0; + return m; +} + +static ncnn::Mat IntArrayMat(int a0, int a1) +{ + ncnn::Mat m(2); + int* p = m; + p[0] = a0; + p[1] = a1; + return m; +} + +static ncnn::Mat IntArrayMat(int a0, int a1, int a2) +{ + ncnn::Mat m(3); + int* p = m; + p[0] = a0; + p[1] = a1; + p[2] = a2; + return m; +} + +static ncnn::Mat IntArrayMat(int a0, int a1, int a2, int a3) +{ + ncnn::Mat m(4); + int* p = m; + p[0] = a0; + p[1] = a1; + p[2] = a2; + p[3] = a3; + return m; +} + +static int test_index_select(const ncnn::Mat& a, const ncnn::Mat& index, int axis) +{ + ncnn::ParamDict pd; + pd.set(0, axis); + + std::vector weights(0); + + std::vector a0(2); + a0[0] = a; + a0[1] = index; + + int ret = test_layer("IndexSelect", pd, weights, a0); + if (ret != 0) + { + fprintf(stderr, "test_index_select failed a.dims=%d a=(%d %d %d %d) index.w=%d axis=%d\n", a.dims, a.w, a.h, a.d, a.c, index.w, axis); + } + + return ret; +} + +static int test_index_select_0() +{ + return 0 + || test_index_select(RandomMat(3, 4, 5, 6), IntArrayMat(1, 0), 0) + || test_index_select(RandomMat(4, 4, 5, 6), IntArrayMat(1, 0), 1) + || test_index_select(RandomMat(3, 4, 7, 6), IntArrayMat(1, 0), 2) + || test_index_select(RandomMat(7, 2, 5, 6), IntArrayMat(1, 0, 3), 3); +} + +static int test_index_select_1() +{ + return 0 + || test_index_select(RandomMat(2, 3, 5), IntArrayMat(1, 0), -3) + || test_index_select(RandomMat(4, 3, 5), IntArrayMat(0, 1), -2) + || test_index_select(RandomMat(6, 4, 5), IntArrayMat(2, 0), -1); +} + +static int test_index_select_2() +{ + return 0 + || test_index_select(RandomMat(8, 6), IntArrayMat(1, 4, 3), 0) + || test_index_select(RandomMat(8, 7), IntArrayMat(3, 5), 1); +} + +static int test_index_select_3() +{ + return 0 + || test_index_select(RandomMat(18), IntArrayMat(1, 7, 9, 15), -1); +} + +int main() +{ + SRAND(7767517); + return 0 + || test_index_select_0() + || test_index_select_1() + || test_index_select_2() + || test_index_select_3(); +} \ No newline at end of file diff --git a/tools/pnnx/tests/ncnn/CMakeLists.txt b/tools/pnnx/tests/ncnn/CMakeLists.txt index 42c3bed32e0..614f17e80d8 100644 --- a/tools/pnnx/tests/ncnn/CMakeLists.txt +++ b/tools/pnnx/tests/ncnn/CMakeLists.txt @@ -203,6 +203,7 @@ pnnx_ncnn_add_test(torch_square) pnnx_ncnn_add_test(torch_tan) pnnx_ncnn_add_test(torch_tanh) pnnx_ncnn_add_test(torch_trunc) +pnnx_ncnn_add_test(torch_index_select) pnnx_ncnn_add_test(convnext_tiny) pnnx_ncnn_add_test(mobilenet_v2) diff --git a/tools/pnnx/tests/ncnn/test_torch_index_select.py b/tools/pnnx/tests/ncnn/test_torch_index_select.py new file mode 100644 index 00000000000..902159cc856 --- /dev/null +++ b/tools/pnnx/tests/ncnn/test_torch_index_select.py @@ -0,0 +1,149 @@ +# Tencent is pleased to support the open source community by making ncnn available. +# +# Copyright (C) 2023 THL A29 Limited, a Tencent company. All rights reserved. +# +# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# https://opensource.org/licenses/BSD-3-Clause +# +# Unless required by applicable law or agreed to in writing, software distributed +# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR +# CONDITIONS OF ANY KIND, either express or implied. See the License for the +# specific language governing permissions and limitations under the License. + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# Tencent is pleased to support the open source community by making ncnn available. +# +# Copyright (C) 2023 THL A29 Limited, a Tencent company. All rights reserved. +# +# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# https://opensource.org/licenses/BSD-3-Clause +# +# Unless required by applicable law or agreed to in writing, software distributed +# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR +# CONDITIONS OF ANY KIND, either express or implied. See the License for the +# specific language governing permissions and limitations under the License. +# Tencent is pleased to support the open source community by making ncnn available. +# +# Copyright (C) 2023 THL A29 Limited, a Tencent company. All rights reserved. +# +# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# https://opensource.org/licenses/BSD-3-Clause +# +# Unless required by applicable law or agreed to in writing, software distributed +# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR +# CONDITIONS OF ANY KIND, either express or implied. See the License for the +# specific language governing permissions and limitations under the License. + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +# 由于ncnn的MemoryData层不支持int64，所以先浮点存档，再转回int64 +# class Model(nn.Module): +# def __init__(self): +# super(Model, self).__init__() + +# def forward(self, x, y, z, d): +# # 1D +# x0 = torch.index_select(x, 0, torch.tensor([7, 9, 11])) +# # 2D +# y0 = torch.index_select(y, 0, torch.tensor([1, 0])) +# y1 = torch.index_select(y, 1, torch.tensor([1, 3, 2])) +# # 3D +# z0 = torch.index_select(z, -3, torch.tensor([1, 0])) +# z1 = torch.index_select(z, -2, torch.tensor([2, 0])) +# z2 = torch.index_select(z, -1, torch.tensor([1, 2])) +# # 4D +# d0 = torch.index_select(d, 0, torch.tensor([1, 0])) +# d1 = torch.index_select(d, 1, torch.tensor([1, 0])) +# d2 = torch.index_select(d, 2, torch.tensor([1, 0])) +# d3 = torch.index_select(d, 3, torch.tensor([1, 0])) + +# return x0, y0, y1, z0, z1, z2, d0, d1, d2, d3 + + +# 成功 +class Model(nn.Module): + def __init__(self): + super(Model, self).__init__() + # 注册所有需要的索引缓冲区 + self.register_buffer("idx_x", torch.tensor([7.0, 9.0, 11.0])) + self.register_buffer("idx_y0", torch.tensor([1.0, 0.0])) + self.register_buffer("idx_y1", torch.tensor([1.0, 3.0, 2.0])) + self.register_buffer("idx_z0", torch.tensor([1.0, 0.0])) + self.register_buffer("idx_z1", torch.tensor([2.0, 0.0])) + self.register_buffer("idx_z2", torch.tensor([1.0, 2.0, 0.0])) + self.register_buffer("idx_d0", torch.tensor([1.0, 0.0, 3.0])) + self.register_buffer("idx_d1", torch.tensor([0.0, 1.0])) + self.register_buffer("idx_d2", torch.tensor([4.0, 3.0, 0.0])) + self.register_buffer("idx_d3", torch.tensor([3.0, 6.0, 2.0])) + + def float_to_int(self, idx_float): + mask = torch.ones_like(idx_float) + return (torch.max(idx_float * mask, mask * 0)).int() + + def forward(self, x, y, z, d): + # 使用辅助函数进行转换 + x0 = torch.index_select(x, 0, self.float_to_int(self.idx_x)) + y0 = torch.index_select(y, 0, self.float_to_int(self.idx_y0)) + y1 = torch.index_select(y, 1, self.float_to_int(self.idx_y1)) + z0 = torch.index_select(z, -3, self.float_to_int(self.idx_z0)) + z1 = torch.index_select(z, -2, self.float_to_int(self.idx_z1)) + z2 = torch.index_select(z, -1, self.float_to_int(self.idx_z2)) + d0 = torch.index_select(d, 0, self.float_to_int(self.idx_d0)) + d1 = torch.index_select(d, 1, self.float_to_int(self.idx_d1)) + d2 = torch.index_select(d, 2, self.float_to_int(self.idx_d2)) + d3 = torch.index_select(d, 3, self.float_to_int(self.idx_d3)) + + return x0, y0, y1, z0, z1, z2, d0, d1, d2, d3 + + +def test(): + net = Model() + net.eval() + + torch.manual_seed(0) + x = torch.rand(36) # 1D + y = torch.rand(5, 7) # 2D + z = torch.rand(3, 5, 8) # 3D + d = torch.rand(4, 3, 6, 7) # 4D + + a = net(x, y, z, d) + + # export torchscript + mod = torch.jit.trace(net, (x, y, z, d)) + mod.save("test_torch_index_select.pt") + + # torchscript to pnnx + import os + + os.system( + "../../src/pnnx test_torch_index_select.pt inputshape=[36],[5,7],[3,5,8],[4,3,6,7]" + ) + + # pnnx inference + import test_torch_index_select_ncnn + + b = test_torch_index_select_ncnn.test_inference() + + for a0, b0 in zip(a, b): + if not torch.allclose(a0, b0, 1e-3, 1e-3): + return False + return True + + +if __name__ == "__main__": + if test(): + exit(0) + else: + exit(1)