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

apps/cpp/convnet.cpp
185 строк5 KB
kolkir
Add tests for Linear layer and update matmul usage
28 мар 2025, 23:53
28 мар 2025, 23:53bc788d8
#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/adam.hpp>
#include <adept/nn/batchnorm2d.hpp>
#include <adept/nn/conv2d.hpp>
#include <adept/nn/cross_entropy.hpp>
#include <adept/nn/linear.hpp>
#include <adept/nn/maxpool2d.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>();
device_t device = device_t::CPU;

index_t epochs = 5;
index_t batch_size = 32;
float32_t lr = 0.01f;
float32_t lr_decay = 0.5;
}  // namespace

class ConvNetImpl : public Module {
 public:
  ConvNetImpl()
      : conv0(Conv2dOptions(1, 64, 7).with_stride(2).with_padding(3).with_bias(false),
              device,
              dtype)  // 14x14x64
        ,
        batch_norm0(64),
        max_pool0(MaxPool2dOptions(3).with_padding(1).with_stride(2))  // 7x7x64
        ,
        conv1(Conv2dOptions(64, 64, 3).with_stride(1).with_padding(1).with_bias(false),
              device,
              dtype)  // 7x7x64
        ,
        batch_norm1(64),
        conv2(Conv2dOptions(64, 256, 7).with_stride(1).with_padding(0).with_bias(false),
              device,
              dtype)  // 1x1x256
        ,
        batch_norm2(256),
        fc(256, 10, device, dtype) {
    register_module("conv0", conv0);
    register_module("batch_norm0", batch_norm0);
    register_module("max_pool0", max_pool0);
    register_module("conv1", conv1);
    register_module("batch_norm1", batch_norm1);
    register_module("conv2", conv2);
    register_module("batch_norm2", batch_norm2);
    register_module("fc", fc);
  }

  Variable forward(Variable x) {
    x = conv0(x);
    x = relu(batch_norm0(x));
    x = max_pool0(x);

    auto identity = x.clone();
    x = conv1(x);
    x = batch_norm1(x);
    x = x + identity;
    x = relu(x);

    x = conv2(x);
    x = batch_norm2(x);

    // view from [N,C,1,1] to [N,C]
    x = x.squeeze(3).squeeze(2);
    x = fc(x);
    return x;
  }

 private:
  Conv2d conv0;
  BarchNorm2d batch_norm0;
  MaxPool2d max_pool0;
  Conv2d conv1;
  BarchNorm2d batch_norm1;
  Conv2d conv2;
  BarchNorm2d batch_norm2;
  Linear fc;
};

ADEPT_MODULE(ConvNet);

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=*/false);
  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=*/false);
  DataLoader test_dataloader(test_dataset, batch_size);

  ConvNet 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);
  }

  Adam 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;
}