Fix small layer issues

include/backend.hpp

@@ -16,6 +16,7 @@ class Backend {
     // Tensor ops
     virtual void print(const CUDANet::Tensor& input) = 0;
     virtual void zero(CUDANet::Tensor& input)        = 0;
+    virtual void fill(CUDANet::Tensor& input, int data) = 0;
 
     virtual void
     copy_to_device(CUDANet::Tensor& tensor, void* data, size_t size) = 0;
@@ -53,7 +54,7 @@ class Backend {
         const CUDANet::Shape out_shape
     ) = 0;
 
-    virtual CUDANet::Tensor& maxPool2d(
+    virtual CUDANet::Tensor& max_pool2d(
         const CUDANet::Tensor& input,
         CUDANet::Tensor& output,
         CUDANet::Shape input_shape,
@@ -63,7 +64,7 @@ class Backend {
         CUDANet::Shape output_shape
     ) = 0;
 
-    virtual CUDANet::Tensor& avgPool2d(
+    virtual CUDANet::Tensor& avg_pool2d(
         const CUDANet::Tensor& input,
         CUDANet::Tensor& output,
         CUDANet::Shape input_shape,
@@ -72,6 +73,29 @@ class Backend {
         CUDANet::Shape padding_shape,
         CUDANet::Shape output_shape
     ) = 0;
+
+    virtual CUDANet::Tensor& batch_norm(
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor& output,
+        CUDANet::Shape input_shape,
+        CUDANet::Tensor& weights,
+        CUDANet::Tensor& biases,
+        CUDANet::Tensor& running_mean,
+        CUDANet::Tensor& running_var,
+        CUDANet::Tensor& epsilon
+    ) = 0;
+
+    virtual CUDANet::Tensor& concat(
+        CUDANet::Tensor& input_a,
+        CUDANet::Tensor& input_b,
+        CUDANet::Tensor& output
+    ) = 0;
+
+    virtual CUDANet::Tensor& add(
+        CUDANet::Tensor& input_a,
+        CUDANet::Tensor& input_b,
+        CUDANet::Tensor& output
+    ) = 0;
 };
 
 }  // namespace CUDANet

include/backend/cuda.cuh

@@ -14,6 +14,7 @@ class CUDA : public Backend {
     // Tensor ops
     void print(const CUDANet::Tensor& input) override;
     void zero(CUDANet::Tensor& input) override;
+    void fill(CUDANet::Tensor &input, int value) override;
     void
     copy_to_device(CUDANet::Tensor& tensor, void* data, size_t size) override;
     void sum(const CUDANet::Tensor& input, CUDANet::Tensor& sum) override;
@@ -49,7 +50,7 @@ class CUDA : public Backend {
         const CUDANet::Shape out_shape
     ) override;
 
-    CUDANet::Tensor& maxPool2d(
+    CUDANet::Tensor& max_pool2d(
         const CUDANet::Tensor& input,
         CUDANet::Tensor& output,
         CUDANet::Shape input_shape,
@@ -59,7 +60,7 @@ class CUDA : public Backend {
         CUDANet::Shape output_shape
     ) override;
 
-    CUDANet::Tensor& avgPool2d(
+    CUDANet::Tensor& avg_pool2d(
         const CUDANet::Tensor& input,
         CUDANet::Tensor& output,
         CUDANet::Shape input_shape,
@@ -67,7 +68,30 @@ class CUDA : public Backend {
         CUDANet::Shape stride_shape,
         CUDANet::Shape padding_shape,
         CUDANet::Shape output_shape
-    ) = 0;
+    ) override;
+
+    CUDANet::Tensor& batch_norm(
+        const CUDANet::Tensor& input,
+        CUDANet::Tensor& output,
+        CUDANet::Shape input_shape,
+        CUDANet::Tensor& weights,
+        CUDANet::Tensor& biases,
+        CUDANet::Tensor& running_mean,
+        CUDANet::Tensor& running_var,
+        CUDANet::Tensor& epsilon
+    ) override;
+
+    CUDANet::Tensor& concat(
+        CUDANet::Tensor& input_a,
+        CUDANet::Tensor& input_b,
+        CUDANet::Tensor& output
+    ) override;
+
+    CUDANet::Tensor& add(
+        CUDANet::Tensor& input_a,
+        CUDANet::Tensor& input_b,
+        CUDANet::Tensor& output
+    ) override;
 };
 
 }  // namespace CUDANet::Backend

include/layers/add.hpp

@@ -1,49 +1,24 @@
-#ifndef CUDANET_ADD_LAYER_H
-#define CUDANET_ADD_LAYER_H
+#pragma once
+
+#include "shape.hpp"
+#include "tensor.hpp"
 
 namespace CUDANet::Layers {
 
 class Add {
   public:
-    /**
-     * @brief Create a new Add layer
-     *
-     * @param inputSize Size of the input arrays
-     */
-    Add(int inputSize);
+    Add(CUDANet::Shape a_shape, CUDANet::Shape b_shape, CUDANet::Backend* backend);
 
-    /**
-     * @brief Destroy the Add layer
-     *
-     */
     ~Add();
 
-    /**
-     * @brief Adds first input to second input
-     *
-     * @param d_inputA Device pointer to the first input
-     * @param d_inputB Device pointer to the second input
-     *
-     */
-    float* forward(const float* inputA, const float* inputB);
+    CUDANet::Tensor&
+    forward(CUDANet::Tensor& input_a, CUDANet::Tensor& input_b);
 
   private:
-    int inputSize;
+    CUDANet::Shape out_shape;
+    CUDANet::Tensor output;
 
-    float* output;
-
-    float* forwardCPU(const float* inputA, const float* inputB);
-
-#ifdef USE_CUDA
-    float* d_output;
-    int gridSize;
-
-    float* forwardCUDA(const float* d_inputA, const float* d_inputB);
-    void initCUDA();
-    void delCUDA();
-#endif
+    CUDANet::Backend *backend;
 };
 
 }  // namespace CUDANet::Layers
-
-#endif  // CUDANET_ADD_LAYER_H

include/layers/batch_norm.hpp

@@ -1,170 +1,54 @@
-#ifndef CUDANET_BATCH_NORM_H
-#define CUDANET_BATCH_NORM_H
+#pragma once
 
-#include <vector>
-
-#include "activation.hpp"
 #include "layer.hpp"
 
 namespace CUDANet::Layers {
 
-class BatchNorm2d : public WeightedLayer, public TwoDLayer {
+class BatchNorm2d : public Layer {
   public:
-    BatchNorm2d(
-        shape2d        inputSize,
-        int            inputChannels,
-        float          epsilon,
-        ActivationType activationType
-    );
+    BatchNorm2d(CUDANet::Shape input_shape, float epsilon, CUDANet::Backend *backend);
 
     ~BatchNorm2d();
 
-    /**
-     * @brief Compute the forward pass of the batchnorm layer
-     *
-     * @param d_input Device pointer to the input
-     * @return float* Device pointer to the output
-     */
-    float* forward(const float* d_input);
+    CUDANet::Tensor& forward(CUDANet::Tensor& input) override;
 
-    /**
-     * @brief Set the weights of the batchnorm layer
-     *
-     * @param weights_input Pointer to the weights
-     */
-    void setWeights(const float* weights_input);
+    CUDANet::Shape input_shape() override;
 
-    /**
-     * @brief Get the weights of the batchnorm layer
-     *
-     * @return std::vector<float>
-     */
-    std::vector<float> getWeights();
+    CUDANet::Shape output_shape() override;
 
-    /**
-     * @brief Set the biases of the batchnorm layer
-     *
-     * @param biases_input Pointer to the biases
-     */
-    void setBiases(const float* biases_input);
+    size_t input_size() override;
 
-    /**
-     * @brief Get the biases of the batchnorm layer
-     *
-     * @return std::vector<float>
-     */
-    std::vector<float> getBiases();
+    size_t output_size() override;
 
-    /**
-     * @brief Set the Running Mean
-     *
-     * @param running_mean_input
-     */
-    void setRunningMean(const float* running_mean_input);
+    void set_weights(void* input) override;
 
-    /**
-     * @brief Get the Running Mean
-     *
-     */
-    std::vector<float> getRunningMean();
+    CUDANet::Tensor& get_weights() override;
 
-    /**
-     * @brief Set the Running Var
-     *
-     * @param running_mean_input
-     */
-    void setRunningVar(const float* running_mean_input);
+    void set_biases(void* input) override;
 
-    /**
-     * @brief Get the Running Var
-     *
-     */
-    std::vector<float> getRunningVar();
+    CUDANet::Tensor& get_biases() override;
 
-    /**
-     * @brief Get output size
-     *
-     * @return int output size
-     */
-    int getOutputSize();
+    void set_running_mean(void* input);
 
-    /**
-     * @brief Get input size
-     *
-     * @return int input size
-     */
-    int getInputSize();
+    CUDANet::Tensor& get_running_mean();
 
-    shape2d getOutputDims();
+    void set_running_var(void* input);
+
+    CUDANet::Tensor& get_running_var();
 
   private:
-    shape2d inputSize;
-    int     inputChannels;
-    float   epsilon;
+    CUDANet::Shape  in_shape;
+    CUDANet::Tensor epsilon;
 
-    int gridSize;
+    CUDANet::Tensor running_mean;
+    CUDANet::Tensor running_var;
 
-#ifdef USE_CUDA
+    CUDANet::Tensor weights;
+    CUDANet::Tensor biases;
 
-    float* d_output;
+    CUDANet::Tensor output;
 
-    float* d_running_mean;
-    float* d_running_var;
-
-    float* d_length;
-    float* d_epsilon;
-
-    float* d_weights;
-    float* d_biases;
-
-    void initCUDA();
-    void delCUDA();
-
-    /**
-     * @brief Copy weights and biases to the device
-     *
-     */
-    void toCuda();
-
-    float* forwardCUDA(const float* d_input);
-
-#endif
-
-    std::vector<float> weights;
-    std::vector<float> biases;
-
-    std::vector<float> running_mean;
-    std::vector<float> running_var;
-
-    Activation* activation;
-
-    float* forwardCPU(const float* input);
-
-    /**
-     * @brief Initialize weights of the batchnorm layer with zeros
-     *
-     */
-    void initializeWeights();
-
-    /**
-     * @brief Initialize biases of the batchnorm layer with zeros
-     *
-     */
-    void initializeBiases();
-
-    /**
-     * @brief Initialize mean of the batchnorm layer with zeros
-     *
-     */
-    void initializeRunningMean();
-
-    /**
-     * @brief Initialize sqrt of variance of the batchnorm layer with ones
-     *
-     */
-    void initializeRunningVar();
+    CUDANet::Backend *backend;
 };
 
 }  // namespace CUDANet::Layers
-
-#endif  // CUDANET_BATCH_NORM_H

include/layers/concat.hpp

@@ -1,5 +1,4 @@
-#ifndef CUDANET_CONCAT_LAYER_H
-#define CUDANET_CONCAT_LAYER_H
+#pragma once
 
 #include "layer.hpp"
 
@@ -11,47 +10,24 @@ namespace CUDANet::Layers {
  */
 class Concat {
   public:
-    /**
-     * @brief Create a new Concat layer
-     *
-     * @param inputASize Size of the first input
-     * @param inputBSize Size of the second input
-     */
-    Concat(const int inputASize, const int inputBSize);
 
-    /**
-     * @brief Destroy the Concat layer
-     *
-     */
+    Concat(const CUDANet::Shape a_shape, const CUDANet::Shape b_shape, CUDANet::Backend *backend);
+
     ~Concat();
 
-    /**
-     * @brief Concatenates the two inputs
-     *
-     * @param d_input_A Device pointer to the first input
-     * @param d_input_B Device pointer to the second input
-     *
-     * @return Device pointer to the output
-     */
-    float* forward(const float* d_input_A, const float* d_input_B);
+    CUDANet::Tensor& forward(CUDANet::Tensor& input_a, CUDANet::Tensor& input_b);
 
-    int getOutputSize();
+    CUDANet::Shape output_shape();
 
   private:
-    int inputASize;
-    int inputBSize;
+    CUDANet::Shape a_shape;
+    CUDANet::Shape b_shape;
 
-    float* forwardCPU(const float* input_A, const float* input_B);
+    CUDANet::Shape out_shape;
+    CUDANet::Tensor output;
 
-#ifdef USE_CUDA
-    float* d_output;
-    float* forwardCUDA(const float* d_input_A, const float* d_input_B);
-
-    void initCUDA();
-    void delCUDA();
-#endif
+    CUDANet::Backend *backend;
 };
 
 }  // namespace CUDANet::Layers
 
-#endif  // CUDANET_CONCAT_LAYER_H

include/layers/conv2d.hpp

@@ -1,8 +1,5 @@
 #pragma once
 
-#include <vector>
-
-#include "activation.hpp"
 #include "layer.hpp"
 
 namespace CUDANet::Layers {

include/layers/dense.hpp

@@ -1,7 +1,5 @@
 #pragma once
 
-#include <vector>
-
 #include "backend.hpp"
 #include "layer.hpp"
 

include/shape.hpp

@@ -21,6 +21,30 @@ class InvalidShapeException : public std::runtime_error {
                   actual
               )
           ) {}
+
+    InvalidShapeException(
+        const std::string& message,
+        const Shape&       shape_a,
+        const Shape&       shape_b
+    )
+        : std::runtime_error(
+              std::format(
+                  "{}. Shape A: [{}], Shape B: [{}]",
+                  message,
+                  format_shape(shape_a),
+                  format_shape(shape_b)
+              )
+          ) {}
+
+  private:
+    static std::string format_shape(const Shape& shape) {
+        std::string result;
+        for (size_t i = 0; i < shape.size(); ++i) {
+            if (i > 0) result += ", ";
+            result += std::to_string(shape[i]);
+        }
+        return result;
+    }
 };
 
 }  // namespace CUDANet

include/tensor.hpp

@@ -45,6 +45,11 @@ public:
 
     void zero();
 
+    template <typename T>
+    void fill(T value) {
+        backend->fill(*this, value);
+    }
+
     template <typename T>
     void set_data(T *data) {
         backend->copy_to_device(*this, data, total_size);

src/backends/cuda/layer_ops.cu

@@ -2,7 +2,7 @@
 #include "kernels/activation_functions.cuh"
 #include "kernels/convolution.cuh"
 #include "kernels/matmul.cuh"
-#include "kernels/pooling.cuh"
+#include "kernels/pool.cuh"
 #include "utils/cuda_helper.cuh"
 
 using namespace CUDANet::Backend;
@@ -98,7 +98,7 @@ CUDANet::Tensor& CUDA::conv2d(
     dim3 grid(
         (out_shape[0] + block.x - 1) / block.x,
         (out_shape[1] + block.y - 1) / block.y,
-        (out_shape[3] + block.z - 1) / block.z
+        (out_shape[2] + block.z - 1) / block.z
     );
 
     Kernels::convolution<<<grid, block>>>(
@@ -112,7 +112,7 @@ CUDANet::Tensor& CUDA::conv2d(
     return output;
 }
 
-CUDANet::Tensor& CUDA::maxPool2d(
+CUDANet::Tensor& CUDA::max_pool2d(
     const CUDANet::Tensor& input,
     CUDANet::Tensor& output,
     CUDANet::Shape input_shape,
@@ -138,7 +138,7 @@ CUDANet::Tensor& CUDA::maxPool2d(
     return output;
 }
 
-CUDANet::Tensor& CUDA::avgPool2d(
+CUDANet::Tensor& CUDA::avg_pool2d(
     const CUDANet::Tensor& input,
     CUDANet::Tensor& output,
     CUDANet::Shape input_shape,
@@ -163,3 +163,91 @@ CUDANet::Tensor& CUDA::avgPool2d(
 
     return output;
 }
+
+CUDANet::Tensor& CUDA::batch_norm(
+    const CUDANet::Tensor& input,
+    CUDANet::Tensor& output,
+    CUDANet::Shape input_shape,
+    CUDANet::Tensor& weights,
+    CUDANet::Tensor& biases,
+    CUDANet::Tensor& running_mean,
+    CUDANet::Tensor& running_var,
+    CUDANet::Tensor& epsilon
+) {
+    auto gridSize =
+        (input_shape[0] * input_shape[1] + BLOCK_SIZE - 1) / BLOCK_SIZE;
+
+
+    for (int i = 0; i < input_shape[2]; i++) {
+        // Subtract mean from input
+        Kernels::vec_scalar_sub<<<gridSize, BLOCK_SIZE>>>(
+            input.data<float>() + i * input_shape[0] * input_shape[1],
+            output.data<float>() + i * input_shape[0] * input_shape[1],
+            &running_mean.data<float>()[i], input_shape[0] * input_shape[1]
+        );
+        CUDA_CHECK(cudaGetLastError());
+
+        // Divide by sqrt(running_var + epsilon)
+        Kernels::vec_scale<<<gridSize, BLOCK_SIZE>>>(
+            output.data<float>() + i * input_shape[0] * input_shape[1],
+            output.data<float>() + i * input_shape[0] * input_shape[1],
+            &running_var.data<float>()[i], epsilon.data<float>(), input_shape[0] * input_shape[1]
+        );
+        CUDA_CHECK(cudaGetLastError());
+
+        // Multiply by weights
+        Kernels::vec_scalar_mul<<<gridSize, BLOCK_SIZE>>>(
+            output.data<float>() + i * input_shape[0] * input_shape[1],
+            output.data<float>() + i * input_shape[0] * input_shape[1], &weights.data<float>()[i],
+            input_shape[0] * input_shape[1]
+        );
+        CUDA_CHECK(cudaGetLastError());
+
+        // Add biases
+        Kernels::vec_scalar_add<<<gridSize, BLOCK_SIZE>>>(
+            output.data<float>() + i * input_shape[0] * input_shape[1],
+            output.data<float>() + i * input_shape[0] * input_shape[1], &biases.data<float>()[i],
+            input_shape[0] * input_shape[1]
+        );
+        CUDA_CHECK(cudaGetLastError());
+    }
+    CUDA_CHECK(cudaDeviceSynchronize());
+    return output;
+}
+
+CUDANet::Tensor& CUDA::concat(
+    CUDANet::Tensor& input_a,
+    CUDANet::Tensor& input_b,
+    CUDANet::Tensor& output
+) {
+    CUDA_CHECK(cudaMemcpy(
+        output.data<float>(), input_a.data<float>(), input_a.size(),
+        cudaMemcpyDeviceToDevice
+    ));
+
+    CUDA_CHECK(cudaMemcpy(
+        output.data<float>() + input_a.numel(), input_b.data<float>(), input_b.size(),
+        cudaMemcpyDeviceToDevice
+    ));
+
+    CUDA_CHECK(cudaGetLastError());
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    return output;
+}
+
+CUDANet::Tensor& CUDA::add(
+    CUDANet::Tensor& input_a,
+    CUDANet::Tensor& input_b,
+    CUDANet::Tensor& output
+) {
+    auto gridSize = (input_a.numel() + BLOCK_SIZE - 1) / BLOCK_SIZE;
+
+    Kernels::vec_vec_add<<<gridSize, BLOCK_SIZE>>>(
+        input_a.data<float>(), input_b.data<float>(), output.data<float>(), input_a.numel()
+    );
+    CUDA_CHECK(cudaGetLastError());
+    CUDA_CHECK(cudaDeviceSynchronize());
+
+    return output;
+}

src/backends/cuda/layers/add.cu

@@ -1,28 +0,0 @@
-#include "add.hpp"
-#include "matmul.cuh"
-#include "cuda_helper.cuh"
-
-using namespace CUDANet::Layers;
-
-void Add::initCUDA() {
-    d_output = nullptr;
-    CUDA_CHECK(cudaMalloc((void**)&d_output, sizeof(float) * inputSize));
-
-    gridSize = (inputSize + BLOCK_SIZE - 1) / BLOCK_SIZE;
-}
-
-void Add::delCUDA() {
-    cudaFree(d_output);
-}
-
-float* Add::forwardCUDA(const float* d_inputA, const float* d_inputB) {
-
-    Kernels::vec_vec_add<<<gridSize, BLOCK_SIZE>>>(
-        d_inputA, d_inputB, d_output, inputSize
-    );
-    CUDA_CHECK(cudaGetLastError());
-    CUDA_CHECK(cudaDeviceSynchronize());
-
-    return d_output;
-
-}

src/backends/cuda/layers/batch_norm.cu

@@ -1,120 +0,0 @@
-#include <vector>
-
-#include "activation.hpp"
-#include "batch_norm.hpp"
-#include "cuda_helper.cuh"
-#include "layer.hpp"
-#include "matmul.cuh"
-#include "vector.cuh"
-
-using namespace CUDANet::Layers;
-
-void BatchNorm2d::initCUDA() {
-    d_output = nullptr;
-    CUDA_CHECK(cudaMalloc(
-        (void **)&d_output,
-        sizeof(float) * inputSize.first * inputSize.second * inputChannels
-    ));
-
-    d_running_mean = nullptr;
-    CUDA_CHECK(
-        cudaMalloc((void **)&d_running_mean, sizeof(float) * inputChannels)
-    );
-
-    d_running_var = nullptr;
-    CUDA_CHECK(
-        cudaMalloc((void **)&d_running_var, sizeof(float) * inputChannels)
-    );
-
-    d_weights = nullptr;
-    CUDA_CHECK(cudaMalloc((void **)&d_weights, sizeof(float) * inputChannels));
-
-    d_biases = nullptr;
-    CUDA_CHECK(cudaMalloc((void **)&d_biases, sizeof(float) * inputChannels));
-
-    d_length     = nullptr;
-    float length = (float)inputSize.first * inputSize.second;
-    CUDA_CHECK(cudaMalloc((void **)&d_length, sizeof(float)));
-    CUDA_CHECK(
-        cudaMemcpy(d_length, &length, sizeof(float), cudaMemcpyHostToDevice)
-    );
-
-    d_epsilon = nullptr;
-    CUDA_CHECK(cudaMalloc((void **)&d_epsilon, sizeof(float)));
-    CUDA_CHECK(
-        cudaMemcpy(d_epsilon, &epsilon, sizeof(float), cudaMemcpyHostToDevice)
-    );
-
-    gridSize =
-        (inputSize.first * inputSize.second + BLOCK_SIZE - 1) / BLOCK_SIZE;
-}
-
-void BatchNorm2d::delCUDA() {
-    cudaFree(d_output);
-    cudaFree(d_running_mean);
-    cudaFree(d_running_var);
-    cudaFree(d_weights);
-    cudaFree(d_biases);
-    cudaFree(d_length);
-    cudaFree(d_epsilon);
-}
-
-void BatchNorm2d::toCuda() {
-    CUDA_CHECK(cudaMemcpy(
-        d_weights, weights.data(), sizeof(float) * inputChannels,
-        cudaMemcpyHostToDevice
-    ));
-    CUDA_CHECK(cudaMemcpy(
-        d_biases, biases.data(), sizeof(float) * inputChannels,
-        cudaMemcpyHostToDevice
-    ));
-    CUDA_CHECK(cudaMemcpy(
-        d_running_mean, running_mean.data(), sizeof(float) * inputChannels,
-        cudaMemcpyHostToDevice
-    ));
-    CUDA_CHECK(cudaMemcpy(
-        d_running_var, running_var.data(), sizeof(float) * inputChannels,
-        cudaMemcpyHostToDevice
-    ));
-}
-
-float *BatchNorm2d::forwardCUDA(const float *d_input) {
-    // Compute per-channel batch normalization
-    for (int i = 0; i < inputChannels; i++) {
-        // Subtract mean from input
-        Kernels::vec_scalar_sub<<<gridSize, BLOCK_SIZE>>>(
-            d_input + i * inputSize.first * inputSize.second,
-            d_output + i * inputSize.first * inputSize.second,
-            &d_running_mean[i], inputSize.first * inputSize.second
-        );
-        CUDA_CHECK(cudaGetLastError());
-
-        // Divide by sqrt(running_var + epsilon)
-        Kernels::vec_scale<<<gridSize, BLOCK_SIZE>>>(
-            d_output + i * inputSize.first * inputSize.second,
-            d_output + i * inputSize.first * inputSize.second,
-            &d_running_var[i], d_epsilon, inputSize.first * inputSize.second
-        );
-        CUDA_CHECK(cudaGetLastError());
-
-        // Multiply by weights
-        Kernels::vec_scalar_mul<<<gridSize, BLOCK_SIZE>>>(
-            d_output + i * inputSize.first * inputSize.second,
-            d_output + i * inputSize.first * inputSize.second, &d_weights[i],
-            inputSize.first * inputSize.second
-        );
-        CUDA_CHECK(cudaGetLastError());
-
-        // Add biases
-        Kernels::vec_scalar_add<<<gridSize, BLOCK_SIZE>>>(
-            d_output + i * inputSize.first * inputSize.second,
-            d_output + i * inputSize.first * inputSize.second, &d_biases[i],
-            inputSize.first * inputSize.second
-        );
-        CUDA_CHECK(cudaGetLastError());
-    }
-
-    activation->activate(d_output);
-
-    return d_output;
-}

src/backends/cuda/layers/concat.cu

@@ -1,31 +0,0 @@
-#include "concat.hpp"
-#include "cuda_helper.cuh"
-
-using namespace CUDANet::Layers;
-
-void Concat::initCUDA() {
-    d_output = nullptr;
-    CUDA_CHECK(
-        cudaMalloc((void**)&d_output, sizeof(float) * (inputASize + inputBSize))
-    );
-}
-
-void Concat::delCUDA() {
-    cudaFree(d_output);
-}
-
-float* Concat::forwardCUDA(const float* d_input_A, const float* d_input_B) {
-    CUDA_CHECK(cudaMemcpy(
-        d_output, d_input_A, sizeof(float) * inputASize,
-        cudaMemcpyDeviceToDevice
-    ));
-
-    CUDA_CHECK(cudaMemcpy(
-        d_output + inputASize, d_input_B, sizeof(float) * inputBSize,
-        cudaMemcpyDeviceToDevice
-    ));
-
-    CUDA_CHECK(cudaDeviceSynchronize());
-
-    return d_output;
-}

src/backends/cuda/tensor_ops.cu

@@ -23,7 +23,12 @@ void CUDA::print(const CUDANet::Tensor &input) {
 }
 
 void CUDA::zero(CUDANet::Tensor &input) {
-    CUDA_CHECK(cudaMemset(input.data<float>(), 0, sizeof(float) * input.numel()));
+    fill(input, 0);
+}
+
+void CUDA::fill(CUDANet::Tensor &input, int value) {
+    CUDA_CHECK(cudaMemset(input.data<float>(), value, sizeof(float) * input.numel()));
+
 }
 
 void CUDA::copy_to_device(CUDANet::Tensor &tensor, void *data, size_t size) {

src/layers/add.cpp

@@ -1,44 +1,28 @@
 #include "add.hpp"
 
-#include <stddef.h>
-
 using namespace CUDANet::Layers;
 
 
-Add::Add(int inputSize)
-    : inputSize(inputSize) {
-
-    output = new float[inputSize];
-
-#ifdef USE_CUDA
-    initCUDA();
-#endif
-    
-}
-
-
-Add::~Add() {
-#ifdef USE_CUDA
-    delCUDA();
-#endif
-}
-
-
-float* Add::forward(const float* inputA, const float* inputB) {
-
-#ifdef USE_CUDA
-    return forwardCUDA(inputA, inputB);
-#else
-    return forwardCPU(inputA, inputB);
-#endif
-
-}
-
-float* Add::forwardCPU(const float* inputA, const float* inputB) {
-    for (size_t i = 0; i < inputSize; i++)
-    {
-        output[i] = inputA[i] + inputB[i];
+Add::Add(CUDANet::Shape a_shape, CUDANet::Shape b_shape, CUDANet::Backend* backend) : backend(backend) {
+    if (a_shape != b_shape) {
+        throw InvalidShapeException(
+            "Add requires matching dimensions", a_shape, b_shape
+        );
     }
 
+    out_shape = a_shape;
+    output = CUDANet::Tensor(out_shape, CUDANet::DType::FLOAT32, backend);
+}
+
+Add::~Add() {}
+
+CUDANet::Tensor&
+Add::forward(CUDANet::Tensor& input_a, CUDANet::Tensor& input_b) {
+    output.zero();
+    backend->add(
+        input_a,
+        input_b,
+        output
+    );
     return output;
 }

src/layers/avg_pooling.cpp

@@ -51,7 +51,7 @@ AvgPool2d::~AvgPool2d() {}
 
 CUDANet::Tensor& AvgPool2d::forward(CUDANet::Tensor& input) {
     output.zero();
-    backend->avgPool2d(
+    backend->avg_pool2d(
         input,
         output,
         in_shape,
@@ -76,7 +76,7 @@ size_t AvgPool2d::input_size() {
 }
 
 size_t AvgPool2d::output_size() {
-    return sizeof(float) * out_shape[0] * out_shape[1] * out_shape[3];
+    return sizeof(float) * out_shape[0] * out_shape[1] * out_shape[2];
 }
 
 void AvgPool2d::set_weights(void* input) {}

src/layers/batch_norm.cpp

@@ -9,125 +9,95 @@
 using namespace CUDANet::Layers;
 
 BatchNorm2d::BatchNorm2d(
-    shape2d        inputSize,
-    int            inputChannels,
-    float          epsilon,
-    ActivationType activationType
+    CUDANet::Shape input_shape,
+    float          eps,
+    CUDANet::Backend *backend
 )
-    : inputSize(inputSize), inputChannels(inputChannels), epsilon(epsilon) {
-    activation = new Activation(
-        activationType, inputSize.first * inputSize.second * inputChannels
+    : in_shape(input_shape), backend(backend)  {
+
+    if (in_shape.size() != 3) {
+        throw InvalidShapeException("input", 3, in_shape.size());
+    }
+
+    epsilon = CUDANet::Tensor({1}, CUDANet::DType::FLOAT32, backend);
+    epsilon.set_data<float>(&eps);
+
+    running_mean = CUDANet::Tensor({in_shape[2]}, CUDANet::DType::FLOAT32, backend);
+    running_mean.zero();
+
+    running_var = CUDANet::Tensor({in_shape[2]}, CUDANet::DType::FLOAT32, backend);
+    running_var.fill(1);
+
+    weights = CUDANet::Tensor({in_shape[2]}, CUDANet::DType::FLOAT32, backend);
+    weights.fill(1);
+
+    biases = CUDANet::Tensor({in_shape[2]}, CUDANet::DType::FLOAT32, backend);
+    biases.zero();
+
+    output = CUDANet::Tensor(in_shape, CUDANet::DType::FLOAT32, backend);
+}
+
+BatchNorm2d::~BatchNorm2d() {}
+
+CUDANet::Tensor& BatchNorm2d::forward(CUDANet::Tensor& input) {
+    output.zero();
+    backend->batch_norm(
+        input,
+        output,
+        in_shape,
+        weights,
+        biases,
+        running_mean,
+        running_var,
+        epsilon
     );
-
-    weights.resize(inputChannels);
-    biases.resize(inputChannels);
-
-    running_mean.resize(inputChannels);
-    running_var.resize(inputChannels);
-
-    initializeWeights();
-    initializeBiases();
-    initializeRunningMean();
-    initializeRunningVar();
-
-#ifdef USE_CUDA
-    initCUDA();
-    toCuda();
-#endif
+    return output;
 }
 
-BatchNorm2d::~BatchNorm2d() {
-#ifdef USE_CUDA
-    delCUDA();
-#endif
+CUDANet::Shape BatchNorm2d::input_shape() {
+    return in_shape;
 }
 
-void BatchNorm2d::initializeWeights() {
-    std::fill(weights.begin(), weights.end(), 1.0f);
+CUDANet::Shape BatchNorm2d::output_shape() {
+    return in_shape;
 }
 
-void BatchNorm2d::initializeBiases() {
-    std::fill(biases.begin(), biases.end(), 0.0f);
+size_t BatchNorm2d::input_size() {
+    return sizeof(float) * in_shape[0] * in_shape[1] * in_shape[2];
 }
 
-void BatchNorm2d::initializeRunningMean() {
-    std::fill(running_mean.begin(), running_mean.end(), 0.0f);
+size_t BatchNorm2d::output_size() {
+    return sizeof(float) * in_shape[0] * in_shape[1] * in_shape[2];
 }
 
-void BatchNorm2d::initializeRunningVar() {
-    std::fill(running_var.begin(), running_var.end(), 1.0f);
+void BatchNorm2d::set_weights(void* input) {
+    weights.set_data<float>(static_cast<float*>(input));
 }
 
-void BatchNorm2d::setWeights(const float* weights_input) {
-    std::copy(weights_input, weights_input + weights.size(), weights.begin());
-#ifdef USE_CUDA
-    toCuda();
-#endif
-}
-
-std::vector<float> BatchNorm2d::getWeights() {
+CUDANet::Tensor& BatchNorm2d::get_weights() {
     return weights;
 }
 
-void BatchNorm2d::setBiases(const float* biases_input) {
-    std::copy(biases_input, biases_input + biases.size(), biases.begin());
-#ifdef USE_CUDA
-    toCuda();
-#endif
+void BatchNorm2d::set_biases(void* input) {
+    biases.set_data<float>(static_cast<float*>(input));
 }
 
-std::vector<float> BatchNorm2d::getBiases() {
+CUDANet::Tensor& BatchNorm2d::get_biases() {
     return biases;
 }
 
-void BatchNorm2d::setRunningMean(const float* running_mean_input) {
-    std::copy(
-        running_mean_input, running_mean_input + inputChannels,
-        running_mean.begin()
-    );
-#ifdef USE_CUDA
-    toCuda();
-#endif
+void BatchNorm2d::set_running_mean(void* input) {
+    running_mean.set_data<float>(static_cast<float*>(input));
 }
 
-std::vector<float> BatchNorm2d::getRunningMean() {
+CUDANet::Tensor& BatchNorm2d::get_running_mean() {
     return running_mean;
 }
 
-void BatchNorm2d::setRunningVar(const float* running_var_input) {
-    std::copy(
-        running_var_input, running_var_input + inputChannels,
-        running_var.begin()
-    );
-#ifdef USE_CUDA
-    toCuda();
-#endif
+void BatchNorm2d::set_running_var(void* input) {
+    running_var.set_data<float>(static_cast<float*>(input));
 }
 
-std::vector<float> BatchNorm2d::getRunningVar() {
+CUDANet::Tensor& BatchNorm2d::get_running_var() {
     return running_var;
 }
-
-int BatchNorm2d::getInputSize() {
-    return inputSize.first * inputSize.second * inputChannels;
-}
-
-int BatchNorm2d::getOutputSize() {
-    return inputSize.first * inputSize.second * inputChannels;
-}
-
-shape2d BatchNorm2d::getOutputDims() {
-    return inputSize;
-}
-
-float* BatchNorm2d::forwardCPU(const float* input) {
-    throw std::logic_error("Not implemented");
-}
-
-float* BatchNorm2d::forward(const float* input) {
-#ifdef USE_CUDA
-    return forwardCUDA(input);
-#else
-    return forwardCPU(input);
-#endif
-}

src/layers/concat.cpp

@@ -1,34 +1,32 @@
-#include <stdexcept>
-
 #include "concat.hpp"
 
 using namespace CUDANet::Layers;
 
-Concat::Concat(const int inputASize, const int inputBSize)
-    : inputASize(inputASize), inputBSize(inputBSize) {
-#ifdef USE_CUDA
-    initCUDA();
-#endif
+Concat::Concat(const CUDANet::Shape a_shape, const CUDANet::Shape b_shape, CUDANet::Backend *backend)
+    : a_shape(a_shape), b_shape(b_shape), backend(backend) {
+    if (a_shape[0] != b_shape[0] || a_shape[1] != b_shape[1]) {
+        throw InvalidShapeException(
+            "Concat requires matching height and width dimensions", a_shape,
+            b_shape
+        );
+    }
+
+    out_shape = {a_shape[0], a_shape[1], a_shape[2] + b_shape[2]};
+    output = CUDANet::Tensor(out_shape, CUDANet::DType::FLOAT32, backend);
 }
 
-Concat::~Concat() {
-#ifdef USE_CUDA
-    delCUDA();
-#endif
+Concat::~Concat() {}
+
+CUDANet::Tensor& Concat::forward(CUDANet::Tensor& input_a, CUDANet::Tensor& input_b) {
+    output.zero();
+    backend->concat(
+        input_a,
+        input_b,
+        output
+    );
+    return output;
 }
 
-float* Concat::forwardCPU(const float* input_A, const float* input_B) {
-    throw std::logic_error("Not implemented");
+CUDANet::Shape Concat::output_shape() {
+    return out_shape;
 }
-
-float* Concat::forward(const float* input_A, const float* input_B) {
-#ifdef USE_CUDA
-    return forwardCUDA(input_A, input_B);
-#else
-    return forwardCPU(input_A, input_B);
-#endif
-}
-
-int Concat::getOutputSize() {
-    return inputASize + inputBSize;
-};

src/layers/conv2d.cpp

@@ -47,7 +47,7 @@ Conv2d::Conv2d(
     };
 
     output = CUDANet::Tensor(
-        Shape{out_shape[0] * out_shape[1] * out_shape[3]},
+        Shape{out_shape[0], out_shape[1], out_shape[2]},
         CUDANet::DType::FLOAT32, backend
     );
 

src/layers/dense.cpp

@@ -30,6 +30,7 @@ Dense::Dense(CUDANet::Shape in_shape, CUDANet::Shape out_shape, CUDANet::Backend
 Dense::~Dense() {}
 
 CUDANet::Tensor& Dense::forward(CUDANet::Tensor& input) {
+    output.zero();
     backend->dense(weights, biases, input, output, in_shape[0], out_shape[0]);
     return output;
 }

src/layers/max_pool.cpp

@@ -41,7 +41,7 @@ MaxPool2d::MaxPool2d(
     };
 
     output = CUDANet::Tensor(
-        Shape{out_shape[0] * out_shape[1] * out_shape[3]},
+        Shape{out_shape[0] * out_shape[1] * out_shape[2]},
         CUDANet::DType::FLOAT32, backend
     );
 }
@@ -50,7 +50,7 @@ MaxPool2d::~MaxPool2d() {}
 
 CUDANet::Tensor& MaxPool2d::forward(CUDANet::Tensor& input) {
     output.zero();
-    backend->maxPool2d(
+    backend->max_pool2d(
         input, output, in_shape, pool_shape, stride_shape, padding_shape,
         out_shape
     );