Add adaptive avg pooling

CMakeLists.txt

@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 3.17)
+cmake_minimum_required(VERSION 3.20)
 
 project(CUDANet
         LANGUAGES CXX CUDA

include/layers/avg_pooling.cuh

@@ -9,11 +9,11 @@ namespace CUDANet::Layers {
 class AvgPooling2d : public SequentialLayer, public TwoDLayer {
   public:
     AvgPooling2d(
-        shape2d           inputSize,
+        shape2d        inputSize,
         int            nChannels,
-        shape2d           poolingSize,
+        shape2d        poolingSize,
-        shape2d           stride,
+        shape2d        stride,
-        shape2d           padding,
+        shape2d        padding,
         ActivationType activationType
     );
     ~AvgPooling2d();
@@ -36,9 +36,9 @@ class AvgPooling2d : public SequentialLayer, public TwoDLayer {
 
     shape2d getOutputDims();
 
-  private:
+  protected:
     shape2d inputSize;
-    int  nChannels;
+    int     nChannels;
     shape2d poolingSize;
     shape2d stride;
     shape2d padding;
@@ -50,6 +50,11 @@ class AvgPooling2d : public SequentialLayer, public TwoDLayer {
     Activation* activation;
 };
 
+class AdaptiveAvgPooling2d : public AvgPooling2d {
+  public:
+    AdaptiveAvgPooling2d(shape2d inputShape, int nChannels, shape2d outputShape, ActivationType activationType);
+};
+
 }  // namespace CUDANet::Layers
 
 #endif  // CUDANET_AVG_POOLING_H

src/layers/avg_pooling.cu

@@ -68,4 +68,24 @@ int AvgPooling2d::getInputSize() {
 
 shape2d AvgPooling2d::getOutputDims() {
     return outputSize;
+}
+
+AdaptiveAvgPooling2d::AdaptiveAvgPooling2d(shape2d inputShape, int nChannels, shape2d outputShape, ActivationType activationType)
+    : AvgPooling2d(inputShape, nChannels, {1, 1}, {1, 1}, {0, 0}, activationType) {
+    
+    stride = {inputShape.first / outputShape.first, inputShape.second / outputShape.second};
+    poolingSize = {
+        inputShape.first - (outputShape.first - 1) * stride.first,
+        inputShape.second - (outputShape.second - 1) * stride.second
+    };
+    padding = {
+        (poolingSize.first - 1) / 2,
+        (poolingSize.second - 1) / 2
+    };
+    outputSize = outputShape;    
+
+    activation = new Activation(activationType, outputSize.first * outputSize.second * nChannels);
+
+    cudaFree(d_output);
+    cudaMalloc((void**)&d_output, sizeof(float) * outputSize.first * outputSize.second * nChannels);
 }

test/layers/test_avg_pooling.cu

@@ -182,4 +182,100 @@ TEST_F(AvgPoolingLayerTest, AvgPoolForwardNonSquarePaddingTest) {
                 0.348f, 0.045f,   0.34075f, 0.43275f, 0.33775f, 0.3955f};
 
     runTest();
-}
+}
+
+
+class AdaptiveAvgPoolingLayerTest : public ::testing::Test {
+  protected:
+    shape2d inputSize;
+    shape2d outputSize;
+    int nChannels;
+    std::vector<float> input;
+    std::vector<float> expected;
+
+    float* d_input;
+    float* d_output;
+    CUDANet::Layers::AdaptiveAvgPooling2d* adaptiveAvgPoolingLayer;
+
+    virtual void SetUp() override {
+        d_input = nullptr;
+        d_output = nullptr;
+        adaptiveAvgPoolingLayer = nullptr;
+    }
+
+    virtual void TearDown() override {
+        cudaFree(d_input);
+    }
+
+    void runTest() {
+        cudaError_t cudaStatus;
+
+        adaptiveAvgPoolingLayer = new CUDANet::Layers::AdaptiveAvgPooling2d(
+            inputSize, nChannels, outputSize, CUDANet::Layers::ActivationType::NONE
+        );
+
+        cudaStatus = cudaMalloc(
+            (void**)&d_input,
+            sizeof(float) * inputSize.first * inputSize.second * nChannels
+        );
+        EXPECT_EQ(cudaStatus, cudaSuccess);
+
+        cudaStatus = cudaMemcpy(
+            d_input, input.data(),
+            sizeof(float) * inputSize.first * inputSize.second * nChannels,
+            cudaMemcpyHostToDevice
+        );
+        EXPECT_EQ(cudaStatus, cudaSuccess);
+
+        d_output = adaptiveAvgPoolingLayer->forward(d_input);
+
+        int outputSize = adaptiveAvgPoolingLayer->getOutputSize();
+
+        std::vector<float> output(outputSize);
+        cudaStatus = cudaMemcpy(
+            output.data(), d_output, sizeof(float) * outputSize,
+            cudaMemcpyDeviceToHost
+        );
+        EXPECT_EQ(cudaStatus, cudaSuccess);
+
+        for (int i = 0; i < output.size(); ++i) {
+            EXPECT_NEAR(expected[i], output[i], 1e-5);
+        }
+
+        delete adaptiveAvgPoolingLayer;
+    }
+};
+
+TEST_F(AdaptiveAvgPoolingLayerTest, AdaptiveAvgPoolForwardTest) {
+    inputSize = {4, 4};
+    outputSize = {2, 2};
+    nChannels = 2;
+
+    input = {
+        // clang-format off
+        // Channel 0
+        0.573f, 0.619f, 0.732f, 0.055f,
+        0.243f, 0.316f, 0.573f, 0.619f,
+        0.712f, 0.055f, 0.243f, 0.316f,
+        0.573f, 0.619f, 0.742f, 0.055f,
+        // Channel 1
+        0.473f, 0.919f, 0.107f, 0.073f,
+        0.073f, 0.362f, 0.973f, 0.059f,
+        0.473f, 0.455f, 0.283f, 0.416f,
+        0.532f, 0.819f, 0.732f, 0.850f
+        // clang-format on
+    };
+
+    expected = {
+        // clang-format off
+        // Channel 0
+        0.43775f, 0.49475f,
+        0.48975f, 0.339f,
+        // Channel 1
+        0.45675f, 0.303f,
+        0.56975f, 0.57025f
+        // clang-format on
+    };
+
+    runTest();
+}

tools/pooling_test.py

@@ -126,6 +126,15 @@ def gen_avg_pool_non_square_padding_test_result():
 
     print_cpp_vector(output)
 
+def gen_adaptive_avg_pool_test_result():
+
+    input = _get_pool_input()
+
+    output = torch.nn.AdaptiveAvgPool2d((2, 2))(input)
+    output = torch.flatten(output)
+
+    print_cpp_vector(output)
+
 
 if __name__ == "__main__":
     print("Generating test results...")
@@ -152,3 +161,8 @@ if __name__ == "__main__":
     gen_avg_pool_non_square_stride_test_result()
     print("Avg pool non square padding test:")
     gen_avg_pool_non_square_padding_test_result()
+
+    print("--------------")
+
+    print("Adaptive avg pool test:")
+    gen_adaptive_avg_pool_test_result()