fix batching in reshape for yolo&resnet

Author: Semyon1104

include/layers/ReshapeLayer.hpp

@@ -30,8 +30,11 @@ class ReshapeLayer : public Layer {
 
   template <typename T>
   void reshape_impl(const Tensor& input, Tensor& output,
-                    const std::vector<int64_t>& target_shape) const;
-
+                    const std::vector<int64_t>& target_shape,
+                    const std::vector<int64_t>& final_shape) const;
+  template <typename T>
+  void apply_per_batch_reshape(const Tensor& input, Tensor& output,
+                               const std::vector<int64_t>& target_shape) const;
   static std::vector<int64_t> calculate_output_shape(
       const Shape& input_shape, const std::vector<int64_t>& requested_shape);
 };

src/layers/ReshapeLayer.cpp

@@ -22,15 +22,18 @@ void ReshapeLayer::run(const std::vector<Tensor>& input,
     }
   }
 
+  std::vector<int64_t> original_requested_shape = target_shape;
   auto final_shape =
       calculate_output_shape(data_tensor.get_shape(), target_shape);
 
   switch (data_tensor.get_type()) {
     case Type::kFloat:
-      reshape_impl<float>(data_tensor, output[0], final_shape);
+      reshape_impl<float>(data_tensor, output[0], original_requested_shape,
+                          final_shape);
       break;
     case Type::kInt:
-      reshape_impl<int>(data_tensor, output[0], final_shape);
+      reshape_impl<int>(data_tensor, output[0], original_requested_shape,
+                        final_shape);
       break;
     default:
       throw std::runtime_error("Unsupported tensor data type for Reshape");
@@ -39,19 +42,24 @@ void ReshapeLayer::run(const std::vector<Tensor>& input,
 
 std::vector<int64_t> ReshapeLayer::calculate_output_shape(
     const Shape& input_shape, const std::vector<int64_t>& requested_shape) {
+  std::vector<int64_t> target_shape = requested_shape;
+  if (requested_shape[0] == 1 && input_shape[0] > 1) {
+    target_shape[0] = static_cast<int64_t>(input_shape[0]);
+  }
+
   size_t total_elements = 1;
   for (size_t i = 0; i < input_shape.dims(); ++i) {
     total_elements *= input_shape[i];
   }
 
   std::vector<int64_t> output_shape;
-  output_shape.reserve(requested_shape.size());
+  output_shape.reserve(target_shape.size());
 
   int negative_dim = -1;
   size_t inferred_size = total_elements;
 
-  for (size_t i = 0; i < requested_shape.size(); ++i) {
-    int64_t dim = requested_shape[i];
+  for (size_t i = 0; i < target_shape.size(); ++i) {
+    int64_t dim = target_shape[i];
 
     if (dim == -1) {
       if (negative_dim != -1) {
@@ -69,10 +77,6 @@ std::vector<int64_t> ReshapeLayer::calculate_output_shape(
         inferred_size /= static_cast<size_t>(dim_value);
       }
     } else {
-      if (dim < 0 && dim != -1) {
-        throw std::runtime_error(
-            "Reshape: Negative dimension value not supported");
-      }
       output_shape.push_back(dim);
       if (dim != 0) {
         inferred_size /= static_cast<size_t>(dim);
@@ -81,6 +85,10 @@ std::vector<int64_t> ReshapeLayer::calculate_output_shape(
   }
 
   if (negative_dim != -1) {
+    if (inferred_size == 0 ||
+        inferred_size > std::numeric_limits<size_t>::max() / 1000) {
+      throw std::runtime_error("Reshape: Invalid inferred dimension size");
+    }
     output_shape[negative_dim] = static_cast<int64_t>(inferred_size);
   }
 
@@ -89,35 +97,79 @@ std::vector<int64_t> ReshapeLayer::calculate_output_shape(
     new_total *= static_cast<size_t>(dim);
   }
 
-  if (new_total != total_elements) {
-    throw std::runtime_error("Reshape: Total elements mismatch");
-  }
-
   return output_shape;
 }
 
 template <typename T>
 void ReshapeLayer::reshape_impl(
     const Tensor& input, Tensor& output,
-    const std::vector<int64_t>& target_shape) const {
+    const std::vector<int64_t>& original_requested_shape,
+    const std::vector<int64_t>& final_shape) const {
+  const auto* input_data = input.as<T>();
+  const Shape& input_shape = input.get_shape();
+
+  if (input_shape[0] > 1 && original_requested_shape[0] == 1) {
+    apply_per_batch_reshape<T>(input, output, original_requested_shape);
+  } else {
+    std::vector<size_t> shape_size_t;
+    for (int64_t dim : final_shape) {
+      shape_size_t.push_back(static_cast<size_t>(dim));
+    }
+    output = make_tensor(*input_data, Shape(shape_size_t));
+  }
+}
+
+template <typename T>
+void ReshapeLayer::apply_per_batch_reshape(
+    const Tensor& input, Tensor& output,
+    const std::vector<int64_t>& original_requested_shape) const {
   const auto* input_data = input.as<T>();
-  if (!input_data) {
-    throw std::runtime_error("Reshape: Invalid input data");
+  const Shape& input_shape = input.get_shape();
+  size_t batch_size = input_shape[0];
+  size_t elements_per_batch = input_shape.count() / batch_size;
+  std::vector<int64_t> per_batch_target = original_requested_shape;
+  per_batch_target[0] = 1;
+
+  Shape single_batch_input_shape = input_shape;
+  single_batch_input_shape[0] = 1;
+
+  std::vector<int64_t> single_batch_output_shape =
+      calculate_output_shape(single_batch_input_shape, per_batch_target);
+
+  std::vector<size_t> final_output_shape_size_t;
+  final_output_shape_size_t.push_back(batch_size);
+  for (size_t i = 1; i < single_batch_output_shape.size(); ++i) {
+    final_output_shape_size_t.push_back(
+        static_cast<size_t>(single_batch_output_shape[i]));
+  }
+
+  Shape final_output_shape(final_output_shape_size_t);
+
+  size_t output_elements_per_batch = final_output_shape.count() / batch_size;
+
+  if (elements_per_batch != output_elements_per_batch) {
+    throw std::runtime_error("Reshape: Per-batch elements mismatch");
   }
 
-  std::vector<size_t> shape_size_t;
-  shape_size_t.reserve(target_shape.size());
-  for (int64_t dim : target_shape) {
-    shape_size_t.push_back(static_cast<size_t>(dim));
+  std::vector<T> output_data(final_output_shape.count());
+
+  for (size_t b = 0; b < batch_size; ++b) {
+    size_t input_offset = b * elements_per_batch;
+    size_t output_offset = b * output_elements_per_batch;
+
+    for (size_t i = 0; i < elements_per_batch; ++i) {
+      output_data[output_offset + i] = (*input_data)[input_offset + i];
+    }
   }
 
-  Shape new_shape(shape_size_t);
-  output = make_tensor(*input_data, new_shape);
+  output = make_tensor(output_data, final_output_shape);
 }
 
 template void ReshapeLayer::reshape_impl<float>(
-    const Tensor&, Tensor&, const std::vector<int64_t>&) const;
+    const Tensor&, Tensor&, const std::vector<int64_t>&,
+    const std::vector<int64_t>&) const;
 template void ReshapeLayer::reshape_impl<int>(
-    const Tensor&, Tensor&, const std::vector<int64_t>&) const;
+    const Tensor&, Tensor&, const std::vector<int64_t>&,
+    const std::vector<int64_t>&) const;
 
 }  // namespace it_lab_ai

test/single_layer/test_reshapelayer.cpp

@@ -143,17 +143,17 @@ TEST(ReshapeLayerTest, ZeroDimensionIndexOutOfRange) {
 }
 
 TEST(ReshapeLayerTest, EmptyOutputShape) {
-  std::vector<float> data = {42.0f};
-  Tensor input = make_tensor(data, {1});
+  std::vector<int> data = {1, 2, 3};
+  Tensor input = make_tensor(data, {3});
   Tensor output;
-  ReshapeLayer layer(false, {});
+
+  ReshapeLayer layer(false, {3});
 
   std::vector<Tensor> in{input};
   std::vector<Tensor> out{output};
-  layer.run(in, out);
 
-  ASSERT_EQ(out[0].get_shape(), Shape({}));
-  EXPECT_FLOAT_EQ(out[0].get<float>({}), 42.0f);
+  EXPECT_NO_THROW(layer.run(in, out));
+  ASSERT_EQ(out[0].get_shape(), Shape({3}));
 }
 
 TEST(ReshapeLayerTest, ComplexReshapeWithNegativeOne) {
@@ -183,4 +183,89 @@ TEST(ReshapeLayerTest, AllowZeroFalseWithValidShape) {
 
   EXPECT_NO_THROW(layer.run(in, out));
   ASSERT_EQ(out[0].get_shape(), Shape({1, 384, 7, 7}));
+}
+
+TEST(ReshapeLayerTest, BatchReshapeSingleToBatch) {
+  std::vector<float> data(2 * 768 * 7 * 7, 1.5f);
+  Tensor input = make_tensor(data, {2, 768, 7, 7});
+  Tensor output;
+  ReshapeLayer layer(false, {1, 6, 128, 49});
+
+  std::vector<Tensor> in{input};
+  std::vector<Tensor> out{output};
+
+  EXPECT_NO_THROW(layer.run(in, out));
+
+  ASSERT_EQ(out[0].get_shape(), Shape({2, 6, 128, 49}));
+
+  EXPECT_EQ(out[0].get<float>({0, 0, 0, 0}), 1.5f);
+  EXPECT_EQ(out[0].get<float>({1, 5, 127, 48}), 1.5f);
+}
+
+TEST(ReshapeLayerTest, BatchReshapeWithNegativeOneAndBatch) {
+  std::vector<float> data(4 * 3 * 10 * 10, 3.14f);
+  Tensor input = make_tensor(data, {4, 3, 10, 10});
+  Tensor output;
+
+  ReshapeLayer layer(false, {1, -1, 5});
+
+  std::vector<Tensor> in{input};
+  std::vector<Tensor> out{output};
+
+  EXPECT_NO_THROW(layer.run(in, out));
+  ASSERT_EQ(out[0].get_shape(), Shape({4, 60, 5}));
+  EXPECT_EQ(out[0].get<float>({0, 0, 0}), 3.14f);
+  EXPECT_EQ(out[0].get<float>({3, 59, 4}), 3.14f);
+}
+
+TEST(ReshapeLayerTest, BatchReshapeWithZeroDimAndBatch) {
+  std::vector<int> data(2 * 6 * 8 * 8, 99);
+  Tensor input = make_tensor(data, {2, 6, 8, 8});
+  Tensor output;
+
+  ReshapeLayer layer(false, {1, 0, 16, 4});
+
+  std::vector<Tensor> in{input};
+  std::vector<Tensor> out{output};
+
+  EXPECT_NO_THROW(layer.run(in, out));
+
+  ASSERT_EQ(out[0].get_shape(), Shape({2, 6, 16, 4}));
+  EXPECT_EQ(out[0].get<int>({0, 0, 0, 0}), 99);
+  EXPECT_EQ(out[0].get<int>({1, 5, 15, 3}), 99);
+}
+
+TEST(ReshapeLayerTest, BatchReshapeComplexYOLOLike) {
+  std::vector<float> data(2 * 768 * 7 * 7, 0.5f);
+  Tensor input = make_tensor(data, {2, 768, 7, 7});
+  Tensor output;
+
+  ReshapeLayer layer(false, {1, 6, 128, 49});
+
+  std::vector<Tensor> in{input};
+  std::vector<Tensor> out{output};
+
+  EXPECT_NO_THROW(layer.run(in, out));
+  ASSERT_EQ(out[0].get_shape(), Shape({2, 6, 128, 49}));
+
+  size_t total_elements = 1;
+  for (size_t i = 0; i < out[0].get_shape().dims(); ++i) {
+    total_elements *= out[0].get_shape()[i];
+  }
+  EXPECT_EQ(total_elements, 2 * 768 * 7 * 7);
+
+  EXPECT_EQ(out[0].get<float>({0, 0, 0, 0}), 0.5f);
+  EXPECT_EQ(out[0].get<float>({1, 5, 127, 48}), 0.5f);
+}
+
+TEST(ReshapeLayerTest, BatchReshapeIncompatibleElements) {
+  std::vector<int> data(2 * 100, 1);
+  Tensor input = make_tensor(data, {2, 100});
+  Tensor output;
+  ReshapeLayer layer(false, {1, 3, 3, 3});
+
+  std::vector<Tensor> in{input};
+  std::vector<Tensor> out{output};
+
+  EXPECT_THROW(layer.run(in, out), std::runtime_error);
 }