TOrchLessCUDA/main.cpp at v2.0 · ANSHAM1/TOrchLessCUDA · 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
116
117
118
119
#include "module.hpp"
#include "optimizer.hpp"

#include "DataLoader.hpp"
#include "MNISTLoader.hpp"

#include <iomanip>


static Sequential buildModel() {
    Sequential model;

    model.Input({ 60000,1,28,28 }, 32);


    model.Conv2D({ 32,1,3,3 }, 1, 1);
    model.ReLU();

    model.Conv2D({ 64,32,3,3 }, 1, 1);
    model.ReLU();

    model.MaxPooling(2, 2);

    model.Flatten();

    model.Dense(128);
    model.ReLU();

    model.Dense(10);

    model.Output("Softmax");


    return model;
}


constexpr int BatchSize = 32;
constexpr int Epochs = 2;


static void TrainingLoop(Sequential& model, ExecutionContext& ctx, Optimizer& opt, Loss& loss, Tensor& images, Tensor& labels) {
    const size_t TotalSamples = images.shape()[0];
    const size_t TotalBatches = (TotalSamples + BatchSize - 1) / BatchSize;

    for (int epoch = 1; epoch <= Epochs; ++epoch) {
        LoadImages imageLoader(images, BatchSize);
        LoadLabels labelLoader(labels, BatchSize);

        float EpochLoss = 0.0f;
        size_t Batch = 0;

        while (imageLoader.hasNext() && labelLoader.hasNext()) {
            Tensor ImagesBatch = imageLoader.next();
            Tensor LabelsBatch = labelLoader.next();

            float LossValue = model.Train(ctx, loss, opt, ImagesBatch, LabelsBatch);

            EpochLoss += LossValue;
            ++Batch;

            const int Progress = static_cast<int>(50.0 * Batch / TotalBatches);

            std::cout << "\rEpoch " << epoch << "/" << Epochs << " [";

            for (int i = 0; i < 50; ++i)
                std::cout << (i < Progress ? '=' : ' ');

            std::cout << "] "
                << std::setw(3)
                << static_cast<int>(100.0 * Batch / TotalBatches)
                << "%";

            std::cout.flush();
        }

        std::cout << "\nEpoch "
            << epoch
            << " Loss: "
            << EpochLoss / Batch
            << "\n\n";
    }
}


int main() {
    auto model = buildModel();

    Tensor Images =
        load_mnist_images(
            "D:/PROJECTs/TorchLess/Dataset/train-images.idx3-ubyte"
        );

    Tensor Labels =
        load_mnist_labels(
            "D:/PROJECTs/TorchLess/Dataset/train-labels.idx1-ubyte"
        );


    ExecutionContext ctx;
    Adam adam(1e-3f);
    Loss loss("CCE");

    TrainingLoop(model, ctx, adam, loss, Images, Labels);


    LoadImages predictionLoader(Images, 5);
    Tensor TestBatch = predictionLoader.next();
    Tensor& Prediction = model.Predict(ctx, TestBatch);

    Prediction.debug_statistics_print("Prediction");

    return 0;
}