wehub-resource-sync
76 lines
1.7 KiB
C++
76 lines
1.7 KiB
C++
//
|
|
// TransformDataset.hpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/11/15.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#ifndef TransformDataset_hpp
|
|
#define TransformDataset_hpp
|
|
|
|
#include <vector>
|
|
#include "Dataset.hpp"
|
|
#include "Example.hpp"
|
|
#include "Transform.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Train {
|
|
|
|
class MNN_PUBLIC BatchTransformDataset : public BatchDataset {
|
|
public:
|
|
BatchTransformDataset(std::shared_ptr<BatchDataset> dataset, std::shared_ptr<BatchTransform> transform) {
|
|
MNN_ASSERT(dataset != nullptr);
|
|
mDataset = dataset;
|
|
mTransform = transform;
|
|
}
|
|
|
|
std::vector<Example> getBatch(std::vector<size_t> indices) override {
|
|
auto batch = mDataset->getBatch(indices);
|
|
if (mTransform != nullptr) {
|
|
batch = mTransform->transformBatch(std::move(batch));
|
|
}
|
|
|
|
return batch;
|
|
}
|
|
|
|
size_t size() override {
|
|
return mDataset->size();
|
|
}
|
|
|
|
private:
|
|
std::shared_ptr<BatchDataset> mDataset;
|
|
std::shared_ptr<BatchTransform> mTransform;
|
|
};
|
|
|
|
class MNN_PUBLIC TransformDataset : public Dataset {
|
|
public:
|
|
TransformDataset(std::shared_ptr<Dataset> dataset, std::shared_ptr<Transform> transform) {
|
|
MNN_ASSERT(dataset != nullptr);
|
|
mDataset = dataset;
|
|
mTransform = transform;
|
|
}
|
|
|
|
Example get(size_t index) override {
|
|
auto example = mDataset->get(index);
|
|
if (mTransform != nullptr) {
|
|
example = mTransform->transformExample(std::move(example));
|
|
}
|
|
|
|
return example;
|
|
}
|
|
|
|
size_t size() override {
|
|
return mDataset->size();
|
|
}
|
|
|
|
private:
|
|
std::shared_ptr<Dataset> mDataset;
|
|
std::shared_ptr<Transform> mTransform;
|
|
};
|
|
|
|
} // namespace Train
|
|
} // namespace MNN
|
|
|
|
#endif
|