wehub-resource-sync
80 lines
2.3 KiB
C++
80 lines
2.3 KiB
C++
//
|
|
// DataLoader.hpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/11/15.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#ifndef DataLoader_hpp
|
|
#define DataLoader_hpp
|
|
|
|
#include <string>
|
|
#include <thread>
|
|
#include <vector>
|
|
#include "BlockingQueue.hpp"
|
|
#include "DataLoaderConfig.hpp"
|
|
#include "Example.hpp"
|
|
namespace MNN {
|
|
namespace Train {
|
|
class BatchDataset;
|
|
class Sampler;
|
|
class BatchTransform;
|
|
class MNN_PUBLIC DataLoader {
|
|
public:
|
|
DataLoader(std::shared_ptr<BatchDataset> dataset, std::shared_ptr<Sampler> sampler,
|
|
std::shared_ptr<DataLoaderConfig> config);
|
|
/*
|
|
When use Windows v141 toolset to compile class having vector of non-copyable element (std::thread, for example),
|
|
copy constructor (or assignment operator) must be deleted explicity, otherwise compile will failed.
|
|
*/
|
|
DataLoader(const DataLoader&) = delete;
|
|
DataLoader& operator = (const DataLoader&) = delete;
|
|
|
|
virtual ~DataLoader() {
|
|
join();
|
|
};
|
|
|
|
void prefetch(size_t nJobs);
|
|
|
|
void workerThread();
|
|
|
|
void join();
|
|
|
|
std::vector<Example> next();
|
|
|
|
void reset();
|
|
|
|
void clean();
|
|
|
|
size_t iterNumber() const;
|
|
size_t size() const;
|
|
static DataLoader* makeDataLoader(std::shared_ptr<BatchDataset> dataset,
|
|
const int batchSize,
|
|
const bool stack = true,
|
|
const bool shuffle = true,
|
|
const int numWorkers = 0);
|
|
static DataLoader* makeDataLoader(std::shared_ptr<BatchDataset> dataset,
|
|
std::vector<std::shared_ptr<BatchTransform>> transforms,
|
|
const int batchSize,
|
|
const bool shuffle = true,
|
|
const int numWorkers = 0);
|
|
|
|
private:
|
|
struct Job {
|
|
std::vector<size_t> job;
|
|
bool quit = false;
|
|
};
|
|
std::shared_ptr<BatchDataset> mDataset;
|
|
std::shared_ptr<Sampler> mSampler;
|
|
std::shared_ptr<DataLoaderConfig> mConfig;
|
|
std::shared_ptr<BlockingQueue<Job>> mJobs;
|
|
std::shared_ptr<BlockingQueue<std::vector<Example>>> mDataQueue;
|
|
std::vector<std::thread> mWorkers;
|
|
};
|
|
|
|
} // namespace Train
|
|
} // namespace MNN
|
|
|
|
#endif // DataLoader_hpp
|