wehub-resource-sync
66 lines
1.5 KiB
C++
66 lines
1.5 KiB
C++
//
|
|
// Dataset.hpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/11/14.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#ifndef Dataset_hpp
|
|
#define Dataset_hpp
|
|
|
|
#include <MNN/MNNDefine.h>
|
|
#include <vector>
|
|
#include "Example.hpp"
|
|
#include "DataLoader.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Train {
|
|
struct MNN_PUBLIC DatasetPtr {
|
|
public:
|
|
std::shared_ptr<BatchDataset> mDataset;
|
|
|
|
DataLoader* createLoader(
|
|
const int batchSize,
|
|
const bool stack = true,
|
|
const bool shuffle = true,
|
|
const int numWorkers = 0);
|
|
~ DatasetPtr() = default;
|
|
template<typename T>
|
|
T* get() const {
|
|
return (T*)mDataset.get();
|
|
}
|
|
};
|
|
|
|
class MNN_PUBLIC BatchDataset {
|
|
public:
|
|
virtual ~BatchDataset() = default;
|
|
|
|
// get batch using given indices
|
|
virtual std::vector<Example> getBatch(std::vector<size_t> indices) = 0;
|
|
|
|
// size of the dataset
|
|
virtual size_t size() = 0;
|
|
};
|
|
|
|
class MNN_PUBLIC Dataset : public BatchDataset {
|
|
public:
|
|
// return a specific example with given index
|
|
virtual Example get(size_t index) = 0;
|
|
|
|
std::vector<Example> getBatch(std::vector<size_t> indices) {
|
|
std::vector<Example> batch;
|
|
batch.reserve(indices.size());
|
|
for (const auto i : indices) {
|
|
batch.emplace_back(get(i));
|
|
}
|
|
MNN_ASSERT(batch.size() != 0);
|
|
return batch;
|
|
}
|
|
};
|
|
|
|
} // namespace Train
|
|
} // namespace MNN
|
|
|
|
#endif /* Dataset_hpp */
|