main.cpp

// MNIST model trainig example
#include <iostream>
#include <torch/csrc/api/include/torch/torch.h>

#ifndef SERIALISED_VERSION
# define USE_C10D_MPI
# include <torch/csrc/distributed/c10d/ProcessGroup.hpp>
# include <torch/csrc/distributed/c10d/ProcessGroupMPI.hpp>
# include <torch/csrc/distributed/c10d/Work.hpp>

# define RUN_MODE_STRING "DISTRIBUTED"
#else
# define RUN_MODE_STRING "SERIALISED "
#endif


#define PROVIDE_TIMING

#if (defined(PROVIDE_TIMING))
# include <chrono>
using namespace std::chrono_literals;
#endif


#define TENSOR_NORMALISE_OP_CST_1 0.1307
#define TENSOR_NORMALISE_OP_CST_2 0.3081


#define CHRONO_SECONDS_TO_DOUBLE(t)                                                                \
    (double)((double)std::chrono::duration_cast<std::chrono::seconds>(t).count() +                 \
             1e-3 * (double)std::chrono::duration_cast<std::chrono::milliseconds>(t % 1s).count())

//


// Define a Convolutional Module
struct Model : torch::nn::Module {
    Model()
        : conv1(torch::nn::Conv2dOptions(1, 10, 5)), conv2(torch::nn::Conv2dOptions(10, 20, 5)),
          fc1(320, 50), fc2(50, 10)
    {
        register_module("conv1", conv1);
        register_module("conv2", conv2);
        register_module("conv2_drop", conv2_drop);
        register_module("fc1", fc1);
        register_module("fc2", fc2);
    }

    torch::Tensor forward(torch::Tensor x)
    {
        x = torch::relu(torch::max_pool2d(conv1->forward(x), 2));
        x = torch::relu(torch::max_pool2d(conv2_drop->forward(conv2->forward(x)), 2));
        x = x.view({-1, 320});
        x = torch::relu(fc1->forward(x));
        x = torch::dropout(x, 0.5, is_training());
        x = fc2->forward(x);
        return torch::log_softmax(x, 1);
    }

    torch::nn::Conv2d    conv1;
    torch::nn::Conv2d    conv2;
    torch::nn::Dropout2d conv2_drop;
    torch::nn::Linear    fc1;
    torch::nn::Linear    fc2;
};


int main(int argc, char *argv[])
{
#ifndef SERIALISED_VERSION
    // Creating MPI Process Group
    auto pg = c10d::ProcessGroupMPI::createProcessGroupMPI();

    // Retrieving MPI environment variables
    const auto numranks = pg->getSize();
    const auto rank     = pg->getRank();
#else
    constexpr auto numranks = 1;
    constexpr auto rank     = 0;
#endif
    if (0 == rank) {
        fprintf(stdout, "\n ****  " RUN_MODE_STRING "  MNIST  toy model  ****\n\n");
    }

    // TRAINING
    // Read train dataset
    const char *kDataRoot     = "./dataset/dataset/";
    auto        train_dataset = torch::data::datasets::MNIST(kDataRoot)
                             .map(torch::data::transforms::Normalize<>(TENSOR_NORMALISE_OP_CST_1,
                                                                       TENSOR_NORMALISE_OP_CST_2))
                             .map(torch::data::transforms::Stack<>());

#ifndef SERIALISED_VERSION
    // Distributed Random Sampler
    auto data_sampler = torch::data::samplers::DistributedRandomSampler(
        train_dataset.size().value(), numranks, rank, false);
#else
    auto data_sampler = torch::data::samplers::RandomSampler(train_dataset.size().value());
#endif

    auto num_train_samples_per_proc = train_dataset.size().value() / numranks;

    // Generate dataloader
    constexpr auto total_batch_size = 64;
    auto           batch_size_per_proc =
        total_batch_size / numranks; // effective batch size in each processor
    auto data_loader =
        torch::data::make_data_loader(std::move(train_dataset), data_sampler, batch_size_per_proc);

    torch::manual_seed(0);

    auto model = std::make_shared<Model>();

    constexpr auto    learning_rate = 1e-2;
    torch::optim::SGD optimizer(model->parameters(), learning_rate);


#if (defined(PROVIDE_TIMING))
    auto hr_clock        = std::chrono::high_resolution_clock();
    auto init_time       = hr_clock.now();
    auto init_train_time = init_time;
    auto end_train_time  = init_train_time;
    auto end_time        = end_train_time;

# ifndef SERIALISED_VERSION
    std::chrono::duration<double> time_allreduce_plus_wait;
    std::chrono::duration<double> time_allreduce;
    std::chrono::duration<double> time_dummy_compute;
    std::chrono::duration<double> time_waitall;
# endif
#endif

    constexpr size_t num_epochs = 10;
    for (size_t epoch = 1; epoch <= num_epochs; ++epoch) {
        size_t num_correct = 0;

        for (auto &batch : *data_loader) {
            auto ip = batch.data;
            auto op = batch.target.squeeze();

            // convert to required formats
            ip = ip.to(torch::kF32);
            op = op.to(torch::kLong);

            // Reset gradients
            model->zero_grad();

            // Execute forward pass
            auto prediction = model->forward(ip);

            auto loss = torch::nll_loss(torch::log_softmax(prediction, 1), op);

            // Backpropagation
            loss.backward();

#ifndef SERIALISED_VERSION
            // Averaging the gradients of the parameters in all the processors
            // Note: This may lag behind DistributedDataParallel (DDP) in performance
            // since this synchronizes parameters after backward pass while DDP
            // overlaps synchronizing parameters and computing gradients in backward pass.

            auto ts = hr_clock.now();

            std::vector<c10::intrusive_ptr<::c10d::Work>> works;
            for (auto &param : model->named_parameters()) {
                std::vector<torch::Tensor> tmp  = {param.value().grad()};
                auto                       ti   = hr_clock.now();
                auto                       work = pg->allreduce(tmp);
                time_allreduce += hr_clock.now() - ti;
                works.push_back(std::move(work));
            }

            auto ti = hr_clock.now();
            for (auto &work : works) {
                try {
                    work->wait();
                } catch (const std::exception &ex) {
                    std::cerr << "Exception received: " << ex.what() << std::endl;
                    pg->abort();
                }
            }
            auto te = hr_clock.now();
            time_waitall += te - ti;
            time_allreduce_plus_wait += te - ts;

            for (auto &param : model->named_parameters()) {
                param.value().grad().data() = param.value().grad().data() / numranks;
            }
#endif

            // Update parameters
            optimizer.step();

            auto guess = prediction.argmax(1);
            num_correct += torch::sum(guess.eq_(op)).item<int64_t>();
        } // end batch loader

        auto accuracy = 100.0 * num_correct / num_train_samples_per_proc;
        std::cout << "Accuracy "
#ifndef SERIALISED_VERSION
                  << "in rank " << rank
#endif
                  << " in epoch " << epoch << " - " << accuracy << std::endl;
    } // end epoch


#if (defined(PROVIDE_TIMING))
    if (0 == rank) {
        end_train_time = hr_clock.now();
    }
#endif

    // TESTING ONLY IN RANK 0
    if (rank == 0) {
        auto test_dataset =
            torch::data::datasets::MNIST(kDataRoot, torch::data::datasets::MNIST::Mode::kTest)
                .map(torch::data::transforms::Normalize<>(TENSOR_NORMALISE_OP_CST_1,
                                                          TENSOR_NORMALISE_OP_CST_2))
                .map(torch::data::transforms::Stack<>());

        auto num_test_samples = test_dataset.size().value();
        auto test_loader = torch::data::make_data_loader(std::move(test_dataset), num_test_samples);

        model->eval(); // enable eval mode to prevent backprop

        size_t num_correct = 0;

        for (auto &batch : *test_loader) {
            auto ip = batch.data;
            auto op = batch.target.squeeze();

            // convert to required format
            ip = ip.to(torch::kF32);
            op = op.to(torch::kLong);

            auto prediction = model->forward(ip);

            auto loss = torch::nll_loss(torch::log_softmax(prediction, 1), op);

            std::cout << "Test loss - " << loss.item<float>() << std::endl;

            auto guess = prediction.argmax(1);

            num_correct += torch::sum(guess.eq_(op)).item<int64_t>();

        } // end test loader

        std::cout << "Num correct - " << num_correct << std::endl;
        std::cout << "Test Accuracy - " << 100.0 * num_correct / num_test_samples << std::endl;
    } // end rank 0


#if (defined(PROVIDE_TIMING))
    if (0 == rank) {
        auto end_train_time = hr_clock.now();

        auto train_time = end_train_time - init_train_time;
        fprintf(stdout, "\nTiming report for  %ld  epochs:\n", num_epochs);
        fprintf(stdout, "\t%s  %7.3lf s\n", "training:", CHRONO_SECONDS_TO_DOUBLE(train_time));
    }
#endif
}