CNN/main.cpp at main · aaryauchamkeri/CNN · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
#include "tensor3d.h"
#include "tensor4d.h"
#include "MNIST_loader.h"
#include "Convolution.h"
#include "ReLU.h"
#include "MaxPool.h"
#include "Flatten.h"
#include "Dense.h"
#include "Softmax.h"

#include <iostream>
#include <cmath>
#include <algorithm>

using namespace std;

Tensor4D unflatten(const vector<float>& flat) {
    Tensor4D out(8, Tensor3D(1, vector<vector<float>>(13, vector<float>(13))));
    int index = 0;
    for (int c = 0; c < 8; c++) {
        for (int i = 0; i < 13; i++) {
            for (int j = 0; j < 13; j++) {
                out[c][0][i][j] = flat[index++];
            }
        }
    }
    return out;
}

int main() {
    Tensor4D train_images;
    vector<int> train_labels;

    MnistLoader::load("../data/train-images-idx3-ubyte", "../data/train-labels-idx1-ubyte", train_images, train_labels);

    int epochs = 10;
    float lr = 0.05;

    Convolution conv({1, 28, 28}, 3, 1, 16, lr);
    ReLU relu;
    MaxPool pool;
    Flatten flatten;
    Dense dense(lr);
    Softmax softmax;


    for (int epoch = 0; epoch < epochs; ++epoch) {
        float total_loss = 0;
        int correct = 0;

        for (int i = 0; i < train_images.size(); ++i) {
            Tensor4D x = {train_images[i]};
            Tensor4D conv_out = conv.forward(x);
            Tensor4D relu_out = relu.forward(conv_out);
            Tensor4D pool_out = pool.forward(relu_out);
            vector<float> flat = flatten.forward(pool_out);
            vector<float> logits = dense.forward(flat);
            vector<float> probs = softmax.forward(logits);

            int label = train_labels[i];
            float loss = -log(probs[label] + 1e-8f);
            total_loss += loss;

            int pred = max_element(probs.begin(), probs.end()) - probs.begin();
            if (pred == label) correct++;
            vector<float> d_softmax = softmax.backward(label);
            vector<float> d_dense = dense.backward(d_softmax);
            Tensor4D d_unflat = unflatten(d_dense);
            Tensor4D d_pool = pool.backward(d_unflat);
            Tensor4D d_relu = relu.backward(d_pool);
            conv.backward(d_relu);
        }

        float avg_loss = total_loss / train_images.size();
        float accuracy = 100.0f * correct / train_images.size();
        cout << "Epoch " << (epoch + 1) << ": Loss = " << avg_loss
             << ", Accuracy = " << accuracy << "%" << endl;
    }

    Tensor4D test_images;
    vector<int> test_labels;

    MnistLoader::load("../data/t10k-images-idx3-ubyte", "../data/t10k-labels-idx1-ubyte", test_images, test_labels);
    float total_loss = 0.0f;
    int correct = 0;
    for (int i = 0; i < test_images.size(); ++i) {
        Tensor4D x = {test_images[i]};
        Tensor4D conv_out = conv.forward(x);
        Tensor4D relu_out = relu.forward(conv_out);
        Tensor4D pool_out = pool.forward(relu_out);
        vector<float> flat = flatten.forward(pool_out);
        vector<float> logits = dense.forward(flat);
        vector<float> probs = softmax.forward(logits);

        int label = test_labels[i];
        float loss = -log(probs[label] + 1e-8f);
        total_loss += loss;

        int pred = max_element(probs.begin(), probs.end()) - probs.begin();
        if (pred == label) correct++;
        vector<float> d_softmax = softmax.backward(label);
        vector<float> d_dense = dense.backward(d_softmax);
        Tensor4D d_unflat = unflatten(d_dense);
        Tensor4D d_pool = pool.backward(d_unflat);
        Tensor4D d_relu = relu.backward(d_pool);
        conv.backward(d_relu);
    }

    float avg_loss = total_loss / test_images.size();
    float accuracy = 100.0f * correct / test_images.size();
    cout << "Test: " << "Loss = " << avg_loss
            << ", Accuracy = " << accuracy << "%" << endl;

    return 0;
}