apps/cpp/mlp.cpp - master - vuron/adept

apps/cpp/mlp.cpp
147 строк4 KB
lexasub
introduce optional usage openblas
09 апр 2025, 11:01
09 апр 2025, 11:013f09613
#include <adept/autograd/autograd.hpp>
#include <adept/data/dataloader.hpp>
#include <adept/data/mnistdataset.hpp>
#include <adept/irange.hpp>
#include <adept/nn/activations.hpp>
#include <adept/nn/cross_entropy.hpp>
#include <adept/nn/linear.hpp>
#include <adept/nn/sgd.hpp>
#include <adept/serialize/fileinput.hpp>
#include <adept/serialize/fileoutput.hpp>

#include <chrono>
#include <filesystem>
#include <iostream>

namespace fs = std::filesystem;
using namespace adept;

namespace {
using DataType = float32_t;
dtype_t dtype = to_dtype<DataType>();
#ifdef WITH_CPU
device_t device = device_t::CPU;
#else
device_t device = device_t::GPU;
#endif
index_t epochs = 5;
index_t batch_size = 32;
float32_t lr = 0.01f;
float32_t lr_decay = 0.5;
}  // namespace

class MLPImpl : public Module {
 public:
  MLPImpl()
      : l1(28 * 28, 512, device, dtype), l2(512, 256, device, dtype), l3(256, 10, device, dtype) {
    register_module("l1", l1);
    register_module("l2", l2);
    register_module("l3", l3);
  }

  Variable forward(Variable input) {
    auto out = l1(input);
    out = relu(out);
    out = l2(out);
    out = relu(out);
    out = l3(out);
    return out;
  }

 private:
  Linear l1;
  Linear l2;
  Linear l3;
};

ADEPT_MODULE(MLP);

int main(int argc, char** argv) {
  if (argc < 2) {
    std::cout << "usage: mlp [path to the MNIST dir] [checkpoint file path]";
    return 0;
  }

  fs::path root = argv[1];
  auto train_images_file = root / "train-images.idx3-ubyte";
  auto train_labels_file = root / "train-labels.idx1-ubyte";

  auto train_dataset =
      std::make_shared<MNISTDataset>(train_images_file, train_labels_file, dtype, device,
                                     /*flat=*/true);
  DataLoader train_dataloader(train_dataset, batch_size);

  auto test_images_file = root / "t10k-images.idx3-ubyte";
  auto test_labels_file = root / "t10k-labels.idx1-ubyte";

  auto test_dataset =
      std::make_shared<MNISTDataset>(test_images_file, test_labels_file, dtype, device,
                                     /*flat=*/true);
  DataLoader test_dataloader(test_dataset, batch_size);

  MLP net;

  // load checkpoint if needed
  fs::path start_checkpoint;
  if (argc == 3) {
    start_checkpoint = argv[2];
    FileInput input(start_checkpoint.native());
    net->load(input);
    input.read("lr", lr);
  }

  SGD optimizer(net->parameters(), lr);
  auto num_batches = train_dataset->size() / batch_size;

  for (auto epoch : irange(epochs)) {
    size_t batch_idx = 0;
    net->train();
    for (auto batch : train_dataloader) {
      Variable x(batch[0], /*requires_grad=*/false);
      Variable y(batch[1], /*requires_grad=*/false);

      auto out = net(x);

      auto loss = cross_entropy_with_logits(out, y);
      if (batch_idx % 64 == 0)
        std::cout << '\r' << "epoch: " << epoch << ", batch: " << batch_idx << "/" << num_batches
                  << ", loss: " << loss.data().at<float32_t>({0, 0}) << std::flush;

      loss.backward();

      optimizer.step();
      optimizer.zero_grad();
      ++batch_idx;
    }
    optimizer.set_lr(optimizer.lr() * lr_decay);

    // save checkpoint
    auto time_point = std::chrono::system_clock::now();
    std::stringstream file_name;
    file_name << "checkpoint_" << epoch << "_" << time_point.time_since_epoch().count();
    FileOutput output(file_name.str());
    net->save(output);
    output.write("lr", optimizer.lr());

    // calculate test loss
    net->eval();
    {
      NoAutoGradGuard guard;
      batch_idx = 0;
      auto total_loss = Tensor::zero({.shape = {1}, .device = device, .dtype = dtype});
      for (auto batch : test_dataloader) {
        Variable x(batch[0]);
        Variable y(batch[1]);

        auto out = net(x);
        auto loss = cross_entropy_with_logits(out, y);
        total_loss += loss.data();
        ++batch_idx;
      }
      std::cout << "\nepoch: " << epoch
                << ", test loss: " << total_loss.at<float32_t>({0}) / batch_idx << std::endl;
    }
  }

  return 0;
}