mnist_dataset_adapter.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <stdexcept>
#include <utility>
 
#include "pulsatrix/dataset.hpp"
#include "pulsatrix/mnist_loader.hpp"
 
namespace pulsatrix {
 
class MnistDatasetAdapter : public Dataset {
public:
    MnistDatasetAdapter(MnistDataset dataset, DeviceBackend* backend)
        : dataset_(std::move(dataset)), backend_(backend) {}
 
    [[nodiscard]] int64_t size() const override { return static_cast<int64_t>(dataset_.images.size()); }
 
    [[nodiscard]] Sample get(int64_t index) const override {
        if (index < 0 || index >= size()) {
            throw std::out_of_range("MnistDatasetAdapter::get: index out of range");
        }
        size_t i = static_cast<size_t>(index);
        Tensor label(Shape({1}), backend_, {static_cast<float>(dataset_.labels[i])});
        return Sample{{dataset_.images[i], std::move(label)}};
    }
 
private:
    MnistDataset dataset_;
    DeviceBackend* backend_;
};
 
}  // namespace pulsatrix