tests/serialize_tests.cpp - with_cpu - vuron/adept

tests/serialize_tests.cpp
94 строки2 KB
kolkir
Make serialization works with mmap
25 янв 2025, 17:25
25 янв 2025, 17:25b021c4e
#include <adept/nn/linear.hpp>
#include <adept/nn/module.hpp>
#include <adept/serialize/fileinput.hpp>
#include <adept/serialize/fileoutput.hpp>

#include <cstdio>

#include "catch.hpp"

using namespace adept;

using DataType = float32_t;
dtype_t dtype = to_dtype<DataType>();
device_t device = device_t::CPU;

class MLPTestImpl : public Module {
 public:
  MLPTestImpl()
      : 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) { return input; }

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

ADEPT_MODULE(MLPTest);

struct SerializeFixture {
  SerializeFixture() { checkpoint_file = tmpnam(nullptr); }

  ~SerializeFixture() noexcept {
    if (!checkpoint_file.empty())
      std::remove(checkpoint_file.c_str());
  }

  std::string checkpoint_file;
};

TEST_CASE_METHOD(SerializeFixture, "Serialize module") {
  REQUIRE_NOTHROW([&]() {
    MLPTest mlp;
    {
      FileOutput output(checkpoint_file);
      mlp->save(output);
    }

    MLPTest mlp2;
    FileInput input(checkpoint_file);
    mlp2->load(input);
  }());
}

TEST_CASE_METHOD(SerializeFixture, "Serialize variable") {
  REQUIRE_NOTHROW([&] {
    auto x0 = Variable(
        Tensor::empty({.shape = {2, 2}, .device = device_t::CPU, .dtype = dtype_t::Float64}));
    auto x1 = Variable(
        Tensor::empty({.shape = {5, 7}, .device = device_t::CPU, .dtype = dtype_t::Int32}));
    auto x2 = Variable(Tensor::from_values(
        {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f}, {2, 2, 4},
        device_t::CPU));
    auto x3 =
        Variable(Tensor::empty({.shape = {3, 3}, .device = device_t::CPU, .dtype = dtype_t::Int8}));

    {
      FileOutput output(checkpoint_file);
      output.write("x_0", x0);
      output.write("x_1", x1);
      output.write("x_2", x2);
      output.write("x_3", x3);
    }
    auto y = Variable(
        Tensor::empty({.shape = {2, 2, 4}, .device = device_t::CPU, .dtype = dtype_t::Float32}));

    FileInput input(checkpoint_file);
    input.read("x_2", y);

    for (index_t b = 0; b < 2; ++b) {
      for (index_t r = 0; r < 2; ++r) {
        for (index_t c = 0; c < 4; ++c) {
          REQUIRE_THAT(x2.data().at<float32_t>({b, r, c}),
                       Catch::WithinRel(y.data().at<float32_t>({b, r, c}), 0.001f));
        }
      }
    }
  }());
}