conv.cpp

#include <iostream>
#include <vector>
#include <cmath>

#include "tensor.h"

namespace tensor {
    Tensor *conv1d(Tensor *input, Tensor *kernel, Tensor *bias, int stride, int padding) {
        std::vector<int> inputShape = input->getShape();
        std::vector<int> kernelShape = kernel->getShape();
        std::vector<int> biasShape = bias->getShape();

        if (inputShape.size() != 3 || kernelShape.size() != 3 || biasShape.size() != 1) {
            throw std::invalid_argument("Invalid input shape");
        }
    
        int batchSize = inputShape[0];      
        int inputChannels = inputShape[1]; 
        int inputWidth = inputShape[2];      
        int kernelChannels = kernelShape[0];
        int outputChannels = kernelShape[1]; 
        int kernelWidth = kernelShape[2];   
    

        if (inputChannels != kernelChannels) {
            throw std::invalid_argument("Input and kernel channels do not match");
        }

        if (biasShape[0] != outputChannels) {
            throw std::invalid_argument("Bias shape does not match output channels");
        }
    
        int paddedWidth = inputWidth + 2 * padding;
        Tensor *paddedInput = Tensor::zeros({batchSize, inputChannels, paddedWidth});
        for (int b = 0; b < batchSize; b++) {
            for (int c = 0; c < inputChannels; c++) {
                for (int w = 0; w < inputWidth; w++) {
                    paddedInput->data[b * inputChannels * paddedWidth + c * paddedWidth + (w + padding)] =
                        input->data[b * inputChannels * inputWidth + c * inputWidth + w];
                }
            }
        }
    

        int outputWidth = (paddedWidth - kernelWidth) / stride + 1;
        std::vector<int> outputShape = {batchSize, outputChannels, outputWidth};
        Tensor *result = Tensor::zeros(outputShape);
    

        for (int b = 0; b < batchSize; b++) {
            for (int oc = 0; oc < outputChannels; oc++) {
                for (int ow = 0; ow < outputWidth; ow++) {
                    float sum = 0.0f;
                    for (int ic = 0; ic < inputChannels; ic++) {
                        for (int kw = 0; kw < kernelWidth; kw++) {
                            int inputIndex = b * inputChannels * paddedWidth + ic * paddedWidth + (ow * stride + kw);
                            int kernelIndex = ic * outputChannels * kernelWidth + oc * kernelWidth + kw;
                            sum += paddedInput->data[inputIndex] * kernel->data[kernelIndex];
                        }
                    }
                    result->data[b * outputChannels * outputWidth + oc * outputWidth + ow] = sum + bias->data[oc];
                }
            }
        }

        if (result->isGradNode()) {
            result->addToPrevious(input);
            result->addToPrevious(kernel);
            result->addToPrevious(bias);
            result->gradFn = [result, input, kernel, bias, stride, padding, inputShape, kernelShape, biasShape, outputWidth, paddedInput]() mutable {
                Tensor *gradInput = Tensor::zeros(inputShape);
                Tensor *gradKernel = Tensor::zeros(kernelShape);
                Tensor *gradBias = Tensor::zeros(biasShape);
    
                int batchSize = inputShape[0];
                int inputChannels = inputShape[1];
                int inputWidth = inputShape[2];
                int outputChannels = kernelShape[1];
                int kernelWidth = kernelShape[2];
                int paddedWidth = inputWidth + 2 * padding;
    
                Tensor *paddedGradInput = Tensor::zeros({batchSize, inputChannels, paddedWidth});
    
                for (int b = 0; b < batchSize; b++) {
                    for (int oc = 0; oc < outputChannels; oc++) {
                        for (int ow = 0; ow < outputWidth; ow++) {
                            float gradOutput = result->grad[b * outputChannels * outputWidth + oc * outputWidth + ow];
                            for (int ic = 0; ic < inputChannels; ic++) {
                                for (int kw = 0; kw < kernelWidth; kw++) {
                                    int inputIndex = b * inputChannels * paddedWidth + ic * paddedWidth + (ow * stride + kw);
                                    int kernelIndex = ic * outputChannels * kernelWidth + oc * kernelWidth + kw;
                                    paddedGradInput->data[inputIndex] += kernel->data[kernelIndex] * gradOutput;
                                    gradKernel->data[kernelIndex] += paddedInput->data[inputIndex] * gradOutput;
                                }
                            }
                            gradBias->data[oc] += gradOutput;
                        }
                    }
                }

                for (int b = 0; b < batchSize; b++) {
                    for (int c = 0; c < inputChannels; c++) {
                        for (int w = 0; w < inputWidth; w++) {
                            gradInput->data[b * inputChannels * inputWidth + c * inputWidth + w] =
                                paddedGradInput->data[b * inputChannels * paddedWidth + c * paddedWidth + (w + padding)];
                        }
                    }
                }
    

                if (input->isGradNode()) {
                    for (int i = 0; i < batchSize * inputChannels * inputWidth; i++) {
                        input->grad[i] += gradInput->data[i];
                    }
                } else {
                    input->grad = gradInput->data;
                }
    
                if (kernel->isGradNode()) {
                    for (int i = 0; i < inputChannels * outputChannels * kernelWidth; i++) {
                        kernel->grad[i] += gradKernel->data[i];
                    }
                } else {
                    kernel->grad = gradKernel->data;
                }
    
                if (bias->isGradNode()) {
                    for (int i = 0; i < outputChannels; i++) {
                        bias->grad[i] += gradBias->data[i];
                    }
                } else {
                    bias->grad = gradBias->data;
                }
            };
        }
    
        return result;
    }

    Tensor* conv2d(Tensor* input, Tensor* kernel, Tensor* bias, int stride_H, int stride_W, int padding_H, int padding_W)
    {
        std::vector<int> inputShape = input->getShape();
        std::vector<int> kernelShape = kernel->getShape();
        std::vector<int> biasShape = bias->getShape();

        if (inputShape.size() != 4 || kernelShape.size() != 4 || biasShape.size() != 1) {
            throw std::invalid_argument("Invalid input shape");
        }
    
        int batchSize = inputShape[0];      
        int inputChannels = inputShape[1]; 
        int inputHeight = inputShape[2];      
        int inputWidth = inputShape[3];      
        int kernelChannels = kernelShape[0];
        int outputChannels = kernelShape[1]; 
        int kernelHeight = kernelShape[2];   
        int kernelWidth = kernelShape[3];    

        if (inputChannels != kernelChannels) {
            throw std::invalid_argument("Input and kernel channels do not match");
        }

        if (biasShape[0] != outputChannels) {
            throw std::invalid_argument("Bias shape does not match output channels");
        }
    
        int paddedHeight = inputHeight + 2 * padding_H;
        int paddedWidth = inputWidth + 2 * padding_W;
        Tensor *paddedInput = Tensor::zeros({batchSize, inputChannels, paddedHeight, paddedWidth});
        for (int b = 0; b < batchSize; b++) {
            for (int c = 0; c < inputChannels; c++) {
                for (int h = 0; h < inputHeight; h++) {
                    for (int w = 0; w < inputWidth; w++) {
                        paddedInput->data[b * inputChannels * paddedHeight * paddedWidth + c * paddedHeight * paddedWidth + (h + padding_H) * paddedWidth + (w + padding_W)] =
                            input->data[b * inputChannels * inputHeight * inputWidth + c * inputHeight * inputWidth + h * inputWidth + w];
                    }
                }
            }
        }

        int outputHeight = (paddedHeight - kernelHeight) / stride_H + 1;
        int outputWidth = (paddedWidth - kernelWidth) / stride_W + 1;
        std::vector<int> outputShape = {batchSize, outputChannels, outputHeight, outputWidth};
        Tensor *result = new Tensor(outputShape);

        for(int batch =0 ; batch < batchSize; batch++)
        {
            for(int channel =0 ; channel < inputChannels; channel++)
            {
                for(int outputChannel =0 ; outputChannel < outputChannels; outputChannel++)
                {
                    for(int outputH =0 ; outputH < outputHeight; outputH++)
                    {
                        for(int outputW =0 ; outputW < outputWidth; outputW++)
                        {
                            float sum = 0.0f;
                            for(int kernelH =0 ; kernelH < kernelHeight; kernelH+=stride_H)
                            {
                                for(int kernelW =0 ; kernelW < kernelWidth; kernelW+=stride_W)
                                {
                                    int inputIndex = batch * input->getStrides()[0] + channel * paddedHeight * paddedWidth + (outputH * stride_H + kernelH) * paddedWidth + (outputW * stride_W + kernelW);
                                    int kernelIndex = channel * outputChannels * kernelHeight * kernelWidth + outputChannel * kernelHeight * kernelWidth + kernelH * kernelWidth + kernelW;
                                    sum += paddedInput->data[inputIndex] * kernel->data[kernelIndex];
                                }
                            }
                            result->data[batch * outputChannels * outputHeight * outputWidth + outputChannel * outputHeight * outputWidth + outputH * outputWidth + outputW] = sum + bias->data[outputChannel];
                        }
                    }
                }
            }
        }

        if(result->isGradNode())
        {
            result->addToPrevious(input);
            result->addToPrevious(kernel);
            result->addToPrevious(bias);
            result->gradFn = [result, input, kernel, bias, stride_H, stride_W, padding_H, padding_W, inputShape, kernelShape, biasShape, outputHeight, outputWidth, paddedInput]() mutable {
                Tensor* gradInput = Tensor::zeros(inputShape);
                Tensor* gradKernel = Tensor::zeros(kernelShape);
                Tensor* gradBias = Tensor::zeros(biasShape);

                int batchSize = inputShape[0];
                int inputChannels = inputShape[1];
                int inputHeight = inputShape[2];
                int inputWidth = inputShape[3];
                int outputChannels = kernelShape[1];
                int kernelHeight = kernelShape[2];
                int kernelWidth = kernelShape[3];
                int paddedHeight = inputHeight + 2 * padding_H;
                int paddedWidth = inputWidth + 2 * padding_W;

                Tensor* paddedGradInput = Tensor::zeros({batchSize, inputChannels, paddedHeight, paddedWidth});

                for(int batch =0 ; batch < batchSize; batch++)
                {
                    for(int outputChannel =0 ; outputChannel < outputChannels; outputChannel++)
                    {
                        for(int outputH =0 ; outputH < outputHeight; outputH++)
                        {
                            for(int outputW =0 ; outputW < outputWidth; outputW++)
                            {
                                float gradOutput = result->grad[batch * outputChannels * outputHeight * outputWidth + outputChannel * outputHeight * outputWidth + outputH * outputWidth + outputW];
                                for(int channel =0 ; channel < inputChannels; channel++)
                                {
                                    for(int kernelH =0 ; kernelH < kernelHeight; kernelH+=stride_H)
                                    {
                                        for(int kernelW =0 ; kernelW < kernelWidth; kernelW+=stride_W)
                                        {
                                            int inputIndex = batch * inputChannels * paddedHeight * paddedWidth + channel * paddedHeight * paddedWidth + (outputH * stride_H + kernelH) * paddedWidth + (outputW * stride_W + kernelW);
                                            int kernelIndex = channel * outputChannels * kernelHeight * kernelWidth + outputChannel * kernelHeight * kernelWidth + kernelH * kernelWidth + kernelW;
                                            paddedGradInput->data[inputIndex] += kernel->data[kernelIndex] * gradOutput;
                                            gradKernel->data[kernelIndex] += paddedInput->data[inputIndex] * gradOutput;
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            };
        }

        return result;

    }
};