vuron/adept: Платформа для обучения и вывода нейронных сетей, с API приближенным к PyTorch.

src/serialize/fileoutput.cpp

77 строк2 KB

kolkir

Make serialization works with mmap

25 янв 2025, 17:25

25 янв 2025, 17:25b021c4e

#include <adept/serialize/fileoutput.hpp>
#include <adept/tensor.hpp>
#include <adept/types_dispatch.hpp>

#include <cstring>
#include <fstream>
#include <type_traits>

namespace adept {

namespace {

void write_len(size_t len, std::ostream& os) {
  os.write(reinterpret_cast<char*>(&len), sizeof(size_t));
}

void write_string(const std::string& str, std::ostream& os) {
  write_len(str.size(), os);
  os.write(str.data(), str.size());
}

template <typename T>
  requires std::is_integral_v<T> || std::is_floating_point_v<T>
void write_number_block(const std::string& name, T value, std::ostream& os) {
  auto block_size = sizeof(size_t) + name.size() + sizeof(T);
  write_len(block_size, os);
  write_string(name, os);
  os.write(reinterpret_cast<char*>(&value), sizeof(T));
}

}  // namespace

FileOutput::FileOutput(const std::string& file_name) {
  out_file_.exceptions(std::ofstream::badbit | std::ofstream::failbit);
  out_file_.open(file_name, std::ios_base::binary);
}

void FileOutput::write(const std::string& name, float32_t value) {
  write_number_block(name, value, out_file_);
}

void FileOutput::write(const std::string& name, float64_t value) {
  write_number_block(name, value, out_file_);
}

void FileOutput::write(const std::string& name, index_t value) {
  write_number_block(name, value, out_file_);
}

void FileOutput::write(const std::string& name, int32_t value) {
  write_number_block(name, value, out_file_);
}

void FileOutput::write(const std::string& name, int8_t value) {
  write_number_block(name, value, out_file_);
}

void FileOutput::write(const std::string& name, const Tensor& tensor) {
  DISPATCH_TYPE(tensor.properties().dtype, [&]() {
    auto numel = tensor.properties().shape.numel();
    auto buffer_size = numel * sizeof(scalar_t);
    auto block_size = sizeof(size_t)                   // blocksize
                      + sizeof(size_t) + name.size()   // name
                      + sizeof(size_t) + buffer_size;  // buffer
    write_len(block_size, out_file_);
    write_string(name, out_file_);
    write_len(buffer_size, out_file_);
    out_file_.write(reinterpret_cast<const char*>(tensor.const_data_ptr<scalar_t>()),
                    numel * sizeof(scalar_t));
  });
}

void FileOutput::write(const std::string& name, const Variable& variable) {
  write(name, variable.data());
};

}  // namespace adept