backprop_conv.cpp

#include <cstdint>
#include <iostream>
#include <vector>
#include <string>
#include <stdio.h>
#include <stdlib.h> 
#include <limits>
#include <chrono>
#include <omp.h>
#include <openacc.h>
#include "backprop_conv.h"
#include "backprop_fcl.h"
#include "consts.h"

using namespace std;

int rotate_180(vector<vector<float>> &input, vector<vector<float>> &output) {
    int H = input.size();
    int W = input[0].size();
    output.resize(H, vector<float>(W));

    for (int i = 0; i < H; i++) {
        for (int j = 0; j < W; j++) {
            output[H - 1 - i][W - 1 - j] = input[i][j];
        }
    }
    return 0;
}


// define convolution 
// define full convolution 

int unflatten(vector<float> &flattened, vector<vector<vector<float>>> &unflattened, int height, int width, int num_kernels){
    unflattened.resize(num_kernels, vector<vector<float>>(height, vector<float>(width)));
    for(int i=0;i<num_kernels;i++){
        for(int j=0;j<height;j++){
            for(int k=0;k<width;k++){
                unflattened[i][j][k] = flattened[i*height*width+j*width+k];
            }
        }
    }
    return 0;
}

// a reverse pool function that can trace back the gradient to the maximum element in the original map 

int reverse_max_pool(vector<vector<vector<float>>> &loss, 
                    vector<vector<vector<float>>> &image_map, 
                    vector<vector<vector<float>>> &unpool, 
                    int stride, int dim){
    unpool.resize(image_map.size(), vector<vector<float>>(image_map[0].size(), vector<float>(image_map[0][0].size(), 0.0f)));
    int row_pool = 0;
    int col_pool = 0;
        
    for(int k=0;k<image_map.size();k++){
        row_pool = 0;
        for(int row_image=0;row_image<image_map[0].size();row_image+=stride){
            col_pool = 0;
            for(int col_image=0;col_image<image_map[0][0].size();col_image+=stride){
                float max_row_index = -1;
                float max_col_index = -1;
                float max = -std::numeric_limits<float>::infinity();
                bool is_max_positive = false;
                for(int i=row_image;i<min(row_image + window, (int)image_map[0].size());i++){
                for(int j=col_image;j<min(col_image +window, (int)image_map[0][0].size());j++){
                    if (image_map[k][i][j] > max){
                        max = image_map[k][i][j];
                        max_row_index = i;
                        max_col_index = j;
                        if (max > 0){
                            is_max_positive = true;
                        }
                    }
                }
            }
            // reverse relu step applied here directly
            if (max_row_index != -1 && max_col_index != -1) {
                unpool[k][max_row_index][max_col_index] = is_max_positive ? loss[k][row_pool][col_pool] : 0.0f;
            }            
            col_pool++;
        }
        row_pool++;
    }
    }
    return 0;
}

// have to change the entire bias matrix now make it a 1d matrix
// lets say i have 16 kernels then only one bias value is used for each kernel so we have 1d matrix of 16 bias values 

int backpropagation_conv(vector<vector<vector<float>>> &loss, 
                        vector<vector<vector<float>>> &image_map, 
                        vector<vector<vector<vector<float>>>> &kernel_list, 
                        vector<float> &bias_accum, 
                        vector<vector<vector<vector<float>>>> &dK_accum,
                        vector<vector<vector<float>>> &loss_inner,
                        vector<vector<vector<vector<float>>>> &rotated_kernel,
                        bool first_itr){
    
    int num_kernels = kernel_list.size(); // number of kernels
    int in_channels = kernel_list[0].size();
    int kH = kernel_list[0][0].size(); // height
    int kW = kernel_list[0][0][0].size(); // width

            
    // Initialize dK
    vector<vector<vector<vector<float>>>> dK(num_kernels, vector<vector<vector<float>>>(in_channels, vector<vector<float>>(kH, vector<float>(kW, 0.0f))));
    //dk/dL is going to be the vector that we try to find using the convolution between the image_map or input and the loss matrices
    chrono::time_point<chrono::system_clock> t1_apply_kernel_backprop = chrono::system_clock::now();
    apply_kernel_backprop(image_map, dK, loss);
    chrono::time_point<chrono::system_clock> t2_apply_kernel_backprop = chrono::system_clock::now();
    chrono::duration<double> elapsed = t2_apply_kernel_backprop - t1_apply_kernel_backprop;
    // cout << "apply_kernel_backprop time: " << elapsed.count() << "s\n";

    chrono::time_point<chrono::system_clock> t1_apply_full_kernel_backprop = chrono::system_clock::now();
    //dX/dL is what we get by rotating kernel by 180 and then doing a full convolution with the loss matrix
    apply_full_kernel_backprop(kernel_list, loss, loss_inner, rotated_kernel, first_itr);
    chrono::time_point<chrono::system_clock> t2_apply_full_kernel_backprop = chrono::system_clock::now();
    chrono::duration<double> elapsed2 = t2_apply_full_kernel_backprop - t1_apply_full_kernel_backprop;
    // cout << "apply_full_kernel_backprop time: " << elapsed2.count() << "s\n";

    chrono::time_point<chrono::system_clock> t1_update_dK = chrono::system_clock::now();
    for (int k = 0; k < num_kernels; ++k) {
        for (int c = 0; c < in_channels; ++c) {
            for (int r = 0; r < kH; ++r) {
                for (int s = 0; s < kW; ++s) {
                    dK_accum[k][c][r][s] += dK[k][c][r][s];
                    // cout << dK_accum[k][c][r][s]<<" ";
                }
            }
        }
    }
    chrono::time_point<chrono::system_clock> t2_update_dK = chrono::system_clock::now();
    chrono::duration<double> elapsed3 = t2_update_dK - t1_update_dK;
    // cout << "update_dK time: " << elapsed3.count() << "s\n";

    // we use one scalar bias for the whole kernel, thats why we have num_kernels number of entries in the bias
    vector<float> dbias(num_kernels, 0.0f);
    int H_out = loss[0].size();
    int W_out = loss[0][0].size();

    for (int k = 0; k < num_kernels; ++k) {
        for (int i = 0; i < H_out; ++i) {
            for (int j = 0; j < W_out; ++j) {
                dbias[k] += loss[k][i][j];
            }
        }
    }
    
    for (int k = 0; k < num_kernels; ++k) {
        bias_accum[k] += dbias[k];
    }

    // for (int k = 0; k < num_kernels; ++k) {
    //     bias[k] -= learning_rate * dbias[k];
    // }

    return 0;
    // here we assume bias is a 1d matrix
    // dbias is just a sum of all the values in one part of a kernel
    // dbias gets updated immediately
}

int apply_kernel_backprop(
    vector<vector<vector<float>>> &image_map,
    vector<vector<vector<vector<float>>>> &dK,
    vector<vector<vector<float>>> &loss) 
{
    int num_kernels = dK.size();             // Output channels
    int input_channels = image_map.size();   // Input channels
    int loss_dim = loss[0].size();           // H_out
    int kernel_size = loss_dim;              // Assuming kernel is same as loss
    int input_height = image_map[0].size();
    int input_width = image_map[0][0].size();
    int center = kernel_size / 2;
    int out_height = input_height - kernel_size + 1;
    int out_width = input_width - kernel_size + 1;
    int kernel_height = dK[0][0].size();
    int kernel_width = dK[0][0][0].size();

    if (loss.size() != num_kernels) {
        cerr << "[ERROR] Mismatch: loss.size() = " << loss.size()
             << ", expected = " << num_kernels << endl;
        exit(1);
    }
    
    for (int k = 0; k < num_kernels; k++) {
        for (int c = 0; c < input_channels; c++) {
            for (int r = 0; r < out_height; r++) {
                for (int s = 0; s < out_width; s++) {
                    for (int i = 0; i < kernel_height; i++) {
                        for (int j = 0; j < kernel_width; j++) {   
                            // if (k == 0 && c == 0 && i == 0 && j == 0 && r == 0 && s == 0) {
                            //     cout << "[Check] loss = " << loss[k][r][s]
                            //          << ", image_map = " << image_map[c][r * stride + i][s * stride + j] << endl;
                            // }                                                     
                            dK[k][c][i][j] += image_map[c][r * stride + i][s * stride + j] * loss[k][r][s];
                        }
                    }
                }
            }
        }
    }
    // cout << "[Debug] dK[1][1][1][0] = " << dK[1][1][1][1] << endl;
    return 0;
}


    int apply_full_kernel_backprop(
        vector<vector<vector<vector<float>>>> &kernel,   // [num_filters][in_channels][kH][kW]
        vector<vector<vector<float>>> &loss,             // [num_filters][H_out][W_out]
        vector<vector<vector<float>>> &dX,
        vector<vector<vector<vector<float>>>> &rotated_kernel,                // [in_channels][H_in][W_in]
        bool first_itr) {
        int num_filters = kernel.size();
        int in_channels = kernel[0].size();
        int kH = kernel[0][0].size();
        int kW = kernel[0][0][0].size();
        int H_out = loss[0].size();
        int W_out = loss[0][0].size();
        int H_in = H_out + kH - 1;
        int W_in = W_out + kW - 1;
    
        // Initialize dX
        dX.clear();
        dX.resize(in_channels, vector<vector<float>>(H_in, vector<float>(W_in, 0.0f)));
    
        // Rotate each kernel
        // if (first_itr) {
        //     for (int k = 0; k < num_filters; ++k) {
        //         for (int c = 0; c < in_channels; ++c) {
        //             rotate_180(kernel[k][c], rotated_kernel[k][c]);
        //             if (rotated_kernel[k][c].size() != kH || rotated_kernel[k][c][0].size() != kW) {
        //                 cerr << "Kernel " << k << ", channel " << c << " has wrong size." << endl;
        //                 exit(1);
        //             }
        //         }
        //     }    
        // }
        // Full convolution
        #pragma omp parallel for collapse(3)
        // #pragma acc parallel loop collapse(3) present(loss, rotated_kernel, dX)
        for (int c = 0; c < in_channels; ++c) {
            for (int i = 0; i < H_in; ++i) {
                for (int j = 0; j < W_in; ++j) {
                    for (int k = 0; k < num_filters; ++k) {
                        for (int m = 0; m < kH; ++m) {
                            for (int n = 0; n < kW; ++n) {
                                int out_i = i - m + kH - 1;
                                int out_j = j - n + kW - 1;
                                if (out_i >= 0 && out_i < H_out && out_j >= 0 && out_j < W_out) {
                                    dX[c][i][j] += loss[k][out_i][out_j] * rotated_kernel[k][c][m][n];
                                }
                            }
                        }
                    }
                }
            }
        }
        return 0;
    }