main.cpp

#include <iostream>
#include <fstream>
#include <iomanip>
#include <string>

#include "tensor.h"
#include "sgd.h"
using namespace std;


static void dump(const string& path, const vector<pair<string, const Tensor*>>& ts, bool grads) {
    ofstream f(path);
    f << scientific << setprecision(17);

    for (const auto& [name, t] : ts) {
        f << name;

        for (int i = 0; i < t->numel(); i++)
            f << " " << (grads ? t->grad[i] : t->data[i]);
        
        f << "\n";
    }
}

int main() {
    // XOR dataset
    Tensor X({4, 2}, {
        0, 0,
        0, 1,
        1, 0,
        1, 1
    });

    Tensor Y({4, 1}, {
        0,
        1,
        1,
        0
    });

    // weights — small random init
    Tensor W1({2, 4}, {
         0.1f, -0.2f,  0.3f,  0.1f,
        -0.1f,  0.2f, -0.3f,  0.2f
    });

    Tensor W2({4, 1}, {
         0.2f,
        -0.1f,
         0.3f,
        -0.2f
    });

    SGD optimizer({&W1, &W2}, 0.01f);

    // training loop
    for (int epoch = 0; epoch < 1000; epoch++) {

        // forward pass
        auto h    = matmul(X, W1);       // (4,2) @ (2,4) = (4,4)
        auto h2   = relu(*h);            // (4,4)
        auto out  = matmul(*h2, W2);     // (4,4) @ (4,1) = (4,1)
        auto loss = mse(*out, Y);        // scalar (1,)

        // print every 100 epochs
        if (epoch % 100 == 0) {
            cout << "epoch " << epoch << "  loss: " << loss->data[0] << endl;
        }

        // backward + update
        Tensor::backward(loss.get());

        if (epoch == 0) dump("grads_step0.txt", {{"W1", &W1}, {"W2", &W2}}, true);

        optimizer.step();
        optimizer.zero_grad();
    }

    dump("weights_final.txt", {{"W1", &W1}, {"W2", &W2}}, false);

    // final predictions
    cout << "\nFinal predictions:" << endl;
    auto h   = matmul(X, W1);
    auto h2  = relu(*h);
    auto out = matmul(*h2, W2);
    
    auto floss = mse(*out, Y);
    cout << "final loss: " << floss->data[0] << endl;

    cout << "Input [0,0] -> " << out->data[0] << "  (expected 0)" << endl;
    cout << "Input [0,1] -> " << out->data[1] << "  (expected 1)" << endl;
    cout << "Input [1,0] -> " << out->data[2] << "  (expected 1)" << endl;
    cout << "Input [1,1] -> " << out->data[3] << "  (expected 0)" << endl;

    return 0;
}