diff --git a/include/layers/SplitLayer.hpp b/include/layers/SplitLayer.hpp new file mode 100644 index 000000000..246c3b5da --- /dev/null +++ b/include/layers/SplitLayer.hpp @@ -0,0 +1,38 @@ +#pragma once +#include +#include +#include + +#include "layers/Layer.hpp" +#include "layers/Tensor.hpp" + +namespace it_lab_ai { + +class SplitLayer : public Layer { + public: + SplitLayer(int axis, std::vector splits) + : axis_(axis), splits_(std::move(splits)) {} + + SplitLayer(int axis, int num_outputs) + : axis_(axis), num_outputs_(num_outputs) {} + void run(const Tensor& input, Tensor& output) override; + void run(const Tensor& input, std::vector& outputs); + + static std::string get_name() { return "SplitLayer"; } + +#ifdef ENABLE_STATISTIC_WEIGHTS + Tensor get_weights() override { return Tensor(); } +#endif + + private: + int axis_; + std::optional> splits_; + std::optional num_outputs_; + + void validate(const Tensor& input) const; + int get_normalized_axis(int rank) const; + template + void split_impl(const Tensor& input, std::vector& outputs) const; +}; + +} // namespace it_lab_ai \ No newline at end of file diff --git a/src/layers/SplitLayer.cpp b/src/layers/SplitLayer.cpp new file mode 100644 index 000000000..3346b08ef --- /dev/null +++ b/src/layers/SplitLayer.cpp @@ -0,0 +1,131 @@ +#include "layers/SplitLayer.hpp" + +#include +#include + +namespace it_lab_ai { + +void SplitLayer::run(const Tensor& input, Tensor& output) { output = input; } + +void SplitLayer::run(const Tensor& input, std::vector& outputs) { + validate(input); + + switch (input.get_type()) { + case Type::kFloat: + split_impl(input, outputs); + break; + case Type::kInt: + split_impl(input, outputs); + break; + default: + throw std::runtime_error("Unsupported tensor data type"); + } +} + +template +void SplitLayer::split_impl(const Tensor& input, + std::vector& outputs) const { + const auto& input_data = *input.as(); + const Shape& shape = input.get_shape(); + const int axis = get_normalized_axis(static_cast(shape.dims())); + + std::vector part_sizes; + if (splits_) { + part_sizes = *splits_; + } else { + const int total_size = static_cast(shape[axis]); + const int base_size = total_size / *num_outputs_; + const int remainder = total_size % *num_outputs_; + + part_sizes.reserve(*num_outputs_); + for (int i = 0; i < *num_outputs_; ++i) { + part_sizes.push_back(i < remainder ? base_size + 1 : base_size); + } + } + + size_t outer_size = 1; + for (int i = 0; i < axis; ++i) { + outer_size *= shape[i]; + } + + size_t inner_size = 1; + for (size_t i = axis + 1; i < shape.dims(); ++i) { + inner_size *= shape[i]; + } + + const size_t input_axis_stride = shape[axis] * inner_size; + + outputs.clear(); + outputs.reserve(part_sizes.size()); + + size_t input_offset = 0; + for (const auto part_size : part_sizes) { + const auto output_axis_size = static_cast(part_size); + + std::vector output_shape_vec(shape.dims()); + for (size_t i = 0; i < shape.dims(); ++i) { + output_shape_vec[i] = + (static_cast(i) == axis) ? output_axis_size : shape[i]; + } + Shape output_shape(output_shape_vec); + + outputs.emplace_back(output_shape, input.get_type()); + auto& output_data = *outputs.back().as(); + + const size_t output_part_size = output_axis_size * inner_size; + + for (size_t outer = 0; outer < outer_size; ++outer) { + const T* input_start = + &input_data[outer * input_axis_stride + input_offset * inner_size]; + T* output_start = &output_data[outer * output_part_size]; + + std::copy_n(input_start, output_part_size, output_start); + } + + input_offset += output_axis_size; + } +} + +void SplitLayer::validate(const Tensor& input) const { + if (input.get_shape().dims() == 0) { + throw std::runtime_error("Cannot split scalar tensor"); + } + + const int axis = + get_normalized_axis(static_cast(input.get_shape().dims())); + const int axis_size = static_cast(input.get_shape()[axis]); + + if (splits_) { + int sum = 0; + for (int s : *splits_) { + if (s <= 0) throw std::runtime_error("Split size must be positive"); + sum += s; + } + if (sum != axis_size) { + throw std::runtime_error("Sum of splits must match axis size"); + } + } else { + if (*num_outputs_ <= 0) { + throw std::runtime_error("num_outputs must be positive"); + } + if (*num_outputs_ > axis_size) { + throw std::runtime_error("num_outputs (" + std::to_string(*num_outputs_) + + ") cannot be greater than axis size (" + + std::to_string(axis_size) + ")"); + } + } +} + +int SplitLayer::get_normalized_axis(int rank) const { + if (axis_ < -rank || axis_ >= rank) { + throw std::runtime_error("Axis out of bounds"); + } + return (axis_ < 0) ? axis_ + rank : axis_; +} + +template void SplitLayer::split_impl(const Tensor&, + std::vector&) const; +template void SplitLayer::split_impl(const Tensor&, + std::vector&) const; + +} // namespace it_lab_ai \ No newline at end of file diff --git a/test/single_layer/test_splitlayer.cpp b/test/single_layer/test_splitlayer.cpp new file mode 100644 index 000000000..93d19d417 --- /dev/null +++ b/test/single_layer/test_splitlayer.cpp @@ -0,0 +1,215 @@ +#include + +#include "gtest/gtest.h" +#include "layers/SplitLayer.hpp" +#include "layers/Tensor.hpp" + +using namespace it_lab_ai; + +TEST(SplitLayerTests, SplitEqualParts1D) { + Tensor input = make_tensor({1, 2, 3, 4, 5, 6}, {6}); + SplitLayer splitter(0, 3); + + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 3); + EXPECT_EQ(outputs[0].get_shape(), Shape({2})); + EXPECT_EQ(outputs[1].get_shape(), Shape({2})); + EXPECT_EQ(outputs[2].get_shape(), Shape({2})); + EXPECT_FLOAT_EQ(outputs[0].get({0}), 1.0f); + EXPECT_FLOAT_EQ(outputs[1].get({0}), 3.0f); + EXPECT_FLOAT_EQ(outputs[2].get({0}), 5.0f); +} + +TEST(SplitLayerTests, SplitVariableParts1D) { + Tensor input = make_tensor({1, 2, 3, 4, 5, 6}, {6}); + SplitLayer splitter(0, {2, 4}); + + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].get_shape(), Shape({2})); + EXPECT_EQ(outputs[1].get_shape(), Shape({4})); + EXPECT_FLOAT_EQ(outputs[0].get({1}), 2.0f); + EXPECT_FLOAT_EQ(outputs[1].get({3}), 6.0f); +} + +TEST(SplitLayerTests, Split2DAlongAxis0) { + Tensor input = make_tensor({1, 2, 3, 4, 5, 6}, {2, 3}); + SplitLayer splitter(0, {1, 1}); + + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].get_shape(), Shape({1, 3})); + EXPECT_EQ(outputs[1].get_shape(), Shape({1, 3})); + EXPECT_FLOAT_EQ(outputs[0].get({0, 2}), 3.0f); + EXPECT_FLOAT_EQ(outputs[1].get({0, 0}), 4.0f); +} + +TEST(SplitLayerTests, Split2DAlongAxis1) { + Tensor input = make_tensor({1, 2, 3, 4, 5, 6}, {2, 3}); + SplitLayer splitter(1, {1, 2}); + + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].get_shape(), Shape({2, 1})); + EXPECT_EQ(outputs[1].get_shape(), Shape({2, 2})); + EXPECT_FLOAT_EQ(outputs[0].get({1, 0}), 4.0f); + EXPECT_FLOAT_EQ(outputs[1].get({0, 1}), 3.0f); +} + +TEST(SplitLayerTests, Split3DEqualParts) { + std::vector data(2 * 3 * 4); + std::iota(data.begin(), data.end(), 0.0f); + Tensor input = make_tensor(data, {2, 3, 4}); + + SplitLayer splitter(1, 3); + + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 3); + EXPECT_EQ(outputs[0].get_shape(), Shape({2, 1, 4})); + EXPECT_EQ(outputs[1].get({1, 0, 3}), 19.0f); +} + +TEST(SplitLayerTests, Split4DVariableParts) { + std::vector data(1 * 3 * 2 * 4); + std::iota(data.begin(), data.end(), 0.0f); + Tensor input = make_tensor(data, {1, 3, 2, 4}); + + SplitLayer splitter(2, {1, 1}); + + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].get_shape(), Shape({1, 3, 1, 4})); + EXPECT_EQ(outputs[1].get({0, 2, 0, 3}), 23.0f); +} + +TEST(SplitLayerTests, SplitNegativeAxis) { + Tensor input = make_tensor({1, 2, 3, 4, 5, 6}, {2, 3}); + SplitLayer splitter(-1, {1, 2}); + + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].get_shape(), Shape({2, 1})); + EXPECT_EQ(outputs[1].get_shape(), Shape({2, 2})); +} + +TEST(SplitLayerTests, InvalidSplitSizes) { + Tensor input = make_tensor({1, 2, 3, 4}, {4}); + + SplitLayer splitter(0, {1, 2}); + + std::vector outputs; + EXPECT_THROW(splitter.run(input, outputs), std::runtime_error); +} + +TEST(SplitLayerTests, EmptyInputTensor) { + Tensor input = make_tensor({}, {0}); + + SplitLayer splitter(0, {}); + + std::vector outputs; + EXPECT_THROW(splitter.run(input, outputs), std::runtime_error); +} + +TEST(SplitLayerTests, Split192IntoTwo96) { + std::vector input_data(1 * 192 * 56 * 56); + std::iota(input_data.begin(), input_data.end(), 0.0f); + Tensor input = make_tensor(input_data, {1, 192, 56, 56}); + + SplitLayer splitter(1, {96, 96}); + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].get_shape(), Shape({1, 96, 56, 56})); + EXPECT_EQ(outputs[1].get_shape(), Shape({1, 96, 56, 56})); + EXPECT_FLOAT_EQ(outputs[0].get({0, 0, 0, 0}), 0.0f); + EXPECT_FLOAT_EQ(outputs[1].get({0, 0, 0, 0}), 96 * 56 * 56); +} + +TEST(SplitLayerTests, UnevenSplitWithRemainder) { + Tensor input = make_tensor({1, 2, 3, 4, 5}, {5}); + SplitLayer splitter(0, 3); + + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 3); + EXPECT_EQ(outputs[0].get_shape(), Shape({2})); + EXPECT_EQ(outputs[1].get_shape(), Shape({2})); + EXPECT_EQ(outputs[2].get_shape(), Shape({1})); + EXPECT_FLOAT_EQ(outputs[0].get({1}), 2.0f); + EXPECT_FLOAT_EQ(outputs[1].get({1}), 4.0f); + EXPECT_FLOAT_EQ(outputs[2].get({0}), 5.0f); +} + +TEST(SplitLayerTests, NumOutputsGreaterThanAxisSize) { + Tensor input = make_tensor({1, 2, 3}, {3}); + SplitLayer splitter(0, 5); + + std::vector outputs; + EXPECT_THROW(splitter.run(input, outputs), std::runtime_error); +} + +TEST(SplitLayerTests, IntegerDataType) { + Tensor input = make_tensor({1, 2, 3, 4, 5, 6}, {2, 3}); + SplitLayer splitter(1, {1, 2}); + + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].get_shape(), Shape({2, 1})); + EXPECT_EQ(outputs[1].get_shape(), Shape({2, 2})); + EXPECT_EQ(outputs[0].get({1, 0}), 4); + EXPECT_EQ(outputs[1].get({0, 1}), 3); +} + +TEST(SplitLayerTests, NegativeAxis2D) { + Tensor input = make_tensor({1, 2, 3, 4}, {2, 2}); + SplitLayer splitter(-2, {1, 1}); + + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].get_shape(), Shape({1, 2})); + EXPECT_EQ(outputs[1].get_shape(), Shape({1, 2})); +} + +TEST(SplitLayerTests, NegativeAxis3D) { + std::vector data(2 * 3 * 4); + std::iota(data.begin(), data.end(), 1.0f); + Tensor input = make_tensor(data, {2, 3, 4}); + + SplitLayer splitter(-1, {1, 3}); + + std::vector outputs; + splitter.run(input, outputs); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].get_shape(), Shape({2, 3, 1})); + EXPECT_EQ(outputs[1].get_shape(), Shape({2, 3, 3})); + EXPECT_FLOAT_EQ(outputs[0].get({1, 2, 0}), 21.0f); +} + +TEST(SplitLayerTests, LargeAxisValue) { + Tensor input = make_tensor({1, 2, 3, 4}, {2, 2}); + + SplitLayer splitter(10, {1, 1}); + std::vector outputs; + EXPECT_THROW(splitter.run(input, outputs), std::runtime_error); +} \ No newline at end of file