src/data/mnistdataset.cpp - master - glebestraikh/adept

src/data/mnistdataset.cpp
111 строк3 KB
kolkir
Massive shaders refacting
27 июн 2025, 23:12
27 июн 2025, 23:12ac55995
#include <adept/data/mnistdataset.hpp>
#include <adept/irange.hpp>
#include <adept/types_dispatch.hpp>
#include <cstring>
#include <fstream>

namespace adept {
namespace {
std::vector<char> read_buffer(const std::string& file_path) {
  std::ifstream file;
  file.open(file_path, std::ios::in | std::ios::binary | std::ios::ate);

  if (!file) {
    THROW_ERROR("Failed to open file {}", file_path);
  }

  auto size = file.tellg();
  std::vector<char> buffer(size);

  file.seekg(0, std::ios::beg);
  file.read(static_cast<char*>(buffer.data()), size);
  file.close();

  return buffer;
}

uint32_t read_header_row(const char* buffer, size_t position) {
  uint32_t value = 0;
  memcpy(&value, buffer + position * sizeof(value), sizeof(value));
  // swap endian
  value = ((value << 8) & 0xFF00FF00) | ((value >> 8) & 0xFF00FF);
  return (value << 16) | (value >> 16);
}
}  // namespace

MNISTDataset::MNISTDataset(const std::string& images_file,
                           const std::string& labels_file,
                           dtype_t dtype,
                           bool flat)
    : tensor_porps_{.shape = {1, 1}, .device = device_t::CPU, .dtype = dtype, .device_id = 0},
      flat_(flat) {
  auto num_images = read_images_file(images_file);
  auto num_labels = read_labels_file(labels_file);
  if (num_images != num_labels) {
    THROW_ERROR("MNISTDataset got different number of images and labels {}, {}", images_file,
                labels_file);
  }
  num_items_ = num_labels;
}

index_t MNISTDataset::read_images_file(const std::string& file_path) {
  images_buffer_ = read_buffer(file_path);
  auto magic = read_header_row(images_buffer_.data(), 0);
  if (magic != 0x00000803) {
    THROW_ERROR("MNISTDataset got invalid images file {}", file_path);
  }
  index_t count = read_header_row(images_buffer_.data(), 1);
  im_rows_ = read_header_row(images_buffer_.data(), 2);
  im_cols_ = read_header_row(images_buffer_.data(), 3);
  im_size_ = im_rows_ * im_cols_;

  if (flat_)
    tensor_porps_.shape = {im_size_};
  else
    tensor_porps_.shape = {1, im_rows_, im_cols_};
  return count;
}

index_t MNISTDataset::read_labels_file(const std::string& file_path) {
  labels_buffer_ = read_buffer(file_path);
  auto magic = read_header_row(labels_buffer_.data(), 0);
  if (magic != 0x00000801) {
    THROW_ERROR("MNISTDataset got invalid labels file {}", file_path);
  }
  index_t count = read_header_row(labels_buffer_.data(), 1);
  return count;
}

index_t MNISTDataset::size() {
  return num_items_;
}

std::vector<Tensor> MNISTDataset::item(index_t idx) {
  if (idx >= num_items_)
    THROW_ERROR("Invalid item index for MNIST dataset");

  std::vector<Tensor> item;
  index_t label_offset = sizeof(uint32_t) + sizeof(uint32_t) + idx;
  auto label = static_cast<uint8_t>(labels_buffer_[label_offset]);

  index_t image_offset = sizeof(uint32_t) + sizeof(uint32_t) * 3 + idx * im_size_;
  char* image_buffer = images_buffer_.data() + image_offset;

  DISPATCH_TYPE(tensor_porps_.dtype, [&]() {
    std::vector<scalar_t> im_data(im_size_);
    for (auto i : irange(im_size_)) {
      auto pixel = static_cast<uint8_t>(image_buffer[i]);
      im_data[i] = static_cast<scalar_t>(pixel) / 255;
    }
    item.emplace_back(Tensor::from_blob(im_data.data(), tensor_porps_));

    int32_t label_value = static_cast<int32_t>(label);
    item.emplace_back(Tensor::from_values({label_value}, {1}, tensor_porps_.device));
  });
  return item;
}

index_t MNISTDataset::item_size() const {
  return 2;
}
}  // namespace adept