main.cpp

#include <torch/torch.h>

#include <algorithm>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

namespace fs = std::filesystem;

// ============================================================
// Activation helpers
// ============================================================

enum class ActType {
    ReLU,
    SELU,
    Tanh,
    Sigmoid
};

std::string act_name(ActType a) {
    switch (a) {
        case ActType::ReLU:    return "relu";
        case ActType::SELU:    return "selu";
        case ActType::Tanh:    return "tanh";
        case ActType::Sigmoid: return "sigmoid";
    }
    return "unknown";
}

torch::Tensor apply_activation(torch::Tensor x, ActType a) {
    switch (a) {
        case ActType::ReLU:    return torch::relu(x);
        case ActType::SELU:    return torch::selu(x);
        case ActType::Tanh:    return torch::tanh(x);
        case ActType::Sigmoid: return torch::sigmoid(x);
    }
    return x;
}

// ============================================================
// CIFAR-10 binary loader
// Expected folder:
// data/cifar-10-batches-bin/
//   data_batch_1.bin ... data_batch_5.bin
//   test_batch.bin
// ============================================================

struct CIFARData {
    torch::Tensor images; // uint8 [N, 3, 32, 32]
    torch::Tensor labels; // int64 [N]
};

CIFARData load_cifar_split(const std::string& folder,
                           const std::vector<std::string>& files) {
    std::vector<uint8_t> image_bytes;
    std::vector<int64_t> labels;

    image_bytes.reserve(files.size() * 10000 * 3072);
    labels.reserve(files.size() * 10000);

    std::vector<char> row(3073);

    for (const auto& file : files) {
        fs::path path = fs::path(folder) / file;
        std::ifstream in(path, std::ios::binary);
        if (!in) {
            throw std::runtime_error("Could not open: " + path.string());
        }

        while (in.read(row.data(), 3073)) {
            labels.push_back(static_cast<unsigned char>(row[0]));
            for (int i = 1; i < 3073; ++i) {
                image_bytes.push_back(static_cast<unsigned char>(row[i]));
            }
        }
    }

    int64_t N = static_cast<int64_t>(labels.size());

    auto images = torch::from_blob(
        image_bytes.data(),
        {N, 3, 32, 32},
        torch::TensorOptions().dtype(torch::kUInt8)
    ).clone();

    auto label_tensor = torch::tensor(labels, torch::TensorOptions().dtype(torch::kInt64));

    return {images, label_tensor};
}

void maybe_take_subset(CIFARData& data, int64_t max_items) {
    if (max_items > 0 && data.images.size(0) > max_items) {
        data.images = data.images.narrow(0, 0, max_items).clone();
        data.labels = data.labels.narrow(0, 0, max_items).clone();
    }
}

// ============================================================
// Basic residual block
// ============================================================

struct BasicBlockImpl : torch::nn::Module {
    ActType act;

    torch::nn::Conv2d conv1{nullptr}, conv2{nullptr}, proj{nullptr};
    torch::nn::BatchNorm2d bn1{nullptr}, bn2{nullptr}, proj_bn{nullptr};
    bool use_proj = false;

    BasicBlockImpl(int64_t in_ch, int64_t out_ch, int64_t stride, ActType act_)
        : act(act_), use_proj(stride != 1 || in_ch != out_ch) {
        conv1 = register_module(
            "conv1",
            torch::nn::Conv2d(
                torch::nn::Conv2dOptions(in_ch, out_ch, 3)
                    .stride(stride)
                    .padding(1)
                    .bias(false)
            )
        );

        bn1 = register_module("bn1", torch::nn::BatchNorm2d(out_ch));

        conv2 = register_module(
            "conv2",
            torch::nn::Conv2d(
                torch::nn::Conv2dOptions(out_ch, out_ch, 3)
                    .stride(1)
                    .padding(1)
                    .bias(false)
            )
        );

        bn2 = register_module("bn2", torch::nn::BatchNorm2d(out_ch));

        if (use_proj) {
            proj = register_module(
                "proj",
                torch::nn::Conv2d(
                    torch::nn::Conv2dOptions(in_ch, out_ch, 1)
                        .stride(stride)
                        .bias(false)
                )
            );
            proj_bn = register_module("proj_bn", torch::nn::BatchNorm2d(out_ch));
        }
    }

    torch::Tensor forward(torch::Tensor x) {
        auto identity = x;

        auto y = conv1->forward(x);
        y = bn1->forward(y);
        y = apply_activation(y, act);

        y = conv2->forward(y);
        y = bn2->forward(y);

        if (use_proj) {
            identity = proj->forward(identity);
            identity = proj_bn->forward(identity);
        }

        y = y + identity;
        y = apply_activation(y, act);
        return y;
    }
};
TORCH_MODULE(BasicBlock);

// ============================================================
// Small CIFAR-style ResNet
// Stem -> 3 residual stages -> global average pool -> FC
// ============================================================

struct ResNetSmallImpl : torch::nn::Module {
    ActType act;
    int64_t in_ch = 16;

    torch::nn::Conv2d stem{nullptr};
    torch::nn::BatchNorm2d stem_bn{nullptr};
    torch::nn::Sequential layer1{nullptr}, layer2{nullptr}, layer3{nullptr};
    torch::nn::AdaptiveAvgPool2d pool{nullptr};
    torch::nn::Linear fc{nullptr};

    explicit ResNetSmallImpl(ActType act_) : act(act_) {
        stem = register_module(
            "stem",
            torch::nn::Conv2d(
                torch::nn::Conv2dOptions(3, 16, 3)
                    .stride(1)
                    .padding(1)
                    .bias(false)
            )
        );

        stem_bn = register_module("stem_bn", torch::nn::BatchNorm2d(16));

        layer1 = register_module("layer1", make_layer(16, 2, 1));
        layer2 = register_module("layer2", make_layer(32, 2, 2));
        layer3 = register_module("layer3", make_layer(64, 2, 2));

        pool = register_module(
            "pool",
            torch::nn::AdaptiveAvgPool2d(torch::nn::AdaptiveAvgPool2dOptions({1, 1}))
        );

        fc = register_module("fc", torch::nn::Linear(64, 10));
    }

    torch::nn::Sequential make_layer(int64_t out_ch, int blocks, int64_t first_stride) {
        torch::nn::Sequential layer;

        layer->push_back(BasicBlock(in_ch, out_ch, first_stride, act));
        in_ch = out_ch;

        for (int i = 1; i < blocks; ++i) {
            layer->push_back(BasicBlock(in_ch, out_ch, 1, act));
        }

        return layer;
    }

    torch::Tensor forward(torch::Tensor x) {
        x = stem->forward(x);
        x = stem_bn->forward(x);
        x = apply_activation(x, act);

        x = layer1->forward(x);
        x = layer2->forward(x);
        x = layer3->forward(x);

        x = pool->forward(x);
        x = x.view({x.size(0), -1});
        x = fc->forward(x);
        return x;
    }
};
TORCH_MODULE(ResNetSmall);

// ============================================================
// Train / eval
// ============================================================

double train_one_epoch(
    ResNetSmall& model,
    const torch::Tensor& train_images,
    const torch::Tensor& train_labels,
    int64_t batch_size,
    torch::optim::SGD& optimizer,
    torch::Device device
) {
    model->train();

    int64_t N = train_images.size(0);
    auto perm = torch::randperm(N, torch::TensorOptions().dtype(torch::kInt64));

    double total_loss = 0.0;

    for (int64_t start = 0; start < N; start += batch_size) {
        int64_t end = std::min(start + batch_size, N);
        auto idx = perm.slice(0, start, end);

        auto x = train_images.index_select(0, idx).to(device).to(torch::kFloat32) / 255.0;
        auto y = train_labels.index_select(0, idx).to(device);

        optimizer.zero_grad();

        auto logits = model->forward(x);
        auto loss = torch::nn::functional::cross_entropy(logits, y);

        loss.backward();
        optimizer.step();

        total_loss += loss.item<double>() * (end - start);
    }

    return total_loss / static_cast<double>(N);
}

std::pair<double, double> evaluate(
    ResNetSmall& model,
    const torch::Tensor& test_images,
    const torch::Tensor& test_labels,
    int64_t batch_size,
    torch::Device device
) {
    model->eval();
    torch::NoGradGuard no_grad;

    int64_t N = test_images.size(0);
    double total_loss = 0.0;
    int64_t correct = 0;

    for (int64_t start = 0; start < N; start += batch_size) {
        int64_t end = std::min(start + batch_size, N);
        auto idx = torch::arange(start, end, torch::TensorOptions().dtype(torch::kInt64));

        auto x = test_images.index_select(0, idx).to(device).to(torch::kFloat32) / 255.0;
        auto y = test_labels.index_select(0, idx).to(device);

        auto logits = model->forward(x);
        auto loss = torch::nn::functional::cross_entropy(logits, y);

        total_loss += loss.item<double>() * (end - start);

        auto pred = logits.argmax(1);
        correct += pred.eq(y).sum().item<int64_t>();
    }

    double avg_loss = total_loss / static_cast<double>(N);
    double acc = static_cast<double>(correct) / static_cast<double>(N);
    return {avg_loss, acc};
}

// ============================================================
// Main
// ============================================================

int main() {
    try {
        torch::manual_seed(42);

        torch::Device device(torch::cuda::is_available() ? torch::kCUDA : torch::kCPU);
        std::cout << "Using device: " << (device.is_cuda() ? "CUDA" : "CPU") << "\n";

        // Adjust this path if needed:
        std::string cifar_dir = "data/cifar-10-batches-bin";

        CIFARData train = load_cifar_split(
            cifar_dir,
            {"data_batch_1.bin", "data_batch_2.bin", "data_batch_3.bin",
             "data_batch_4.bin", "data_batch_5.bin"}
        );

        CIFARData test = load_cifar_split(
            cifar_dir,
            {"test_batch.bin"}
        );

        // Lower these if you want faster experiments.
        int64_t max_train = 50000; // e.g. 20000 for quick testing
        int64_t max_test  = 10000; // e.g. 2000 for quick testing

        maybe_take_subset(train, max_train);
        maybe_take_subset(test, max_test);

        std::cout << "Train samples: " << train.images.size(0) << "\n";
        std::cout << "Test samples : " << test.images.size(0) << "\n";

        const int epochs = 10;
        const int64_t batch_size = 128;
        const double lr = 0.1;
        const double momentum = 0.9;
        const double weight_decay = 1e-4;

        std::vector<ActType> activations = {
            ActType::ReLU,
            ActType::SELU,
            ActType::Tanh,
            ActType::Sigmoid
        };

        std::ofstream csv("resnet_cifar10_activation_compare.csv");
        csv << "activation,epoch,train_loss,test_loss,test_accuracy\n";

        std::cout << std::fixed << std::setprecision(4);

        for (auto act : activations) {
            torch::manual_seed(42);

            ResNetSmall model(act);
            model->to(device);

            torch::optim::SGD optimizer(
                model->parameters(),
                torch::optim::SGDOptions(lr).momentum(momentum).weight_decay(weight_decay)
            );

            std::cout << "\n=== Activation: " << act_name(act) << " ===\n";

            for (int epoch = 1; epoch <= epochs; ++epoch) {
                double train_loss = train_one_epoch(
                    model, train.images, train.labels, batch_size, optimizer, device
                );

                auto [test_loss, test_acc] = evaluate(
                    model, test.images, test.labels, batch_size, device
                );

                std::cout
                    << "epoch " << epoch
                    << " | train_loss = " << train_loss
                    << " | test_loss = " << test_loss
                    << " | test_acc = " << test_acc
                    << "\n";

                csv << act_name(act) << ","
                    << epoch << ","
                    << train_loss << ","
                    << test_loss << ","
                    << test_acc << "\n";
            }
        }

        csv.close();
        std::cout << "\nSaved: resnet_cifar10_activation_compare.csv\n";
    }
    catch (const std::exception& e) {
        std::cerr << "ERROR: " << e.what() << "\n";
        return 1;
    }

    return 0;
}