MNN/tools/train/source/demo/dataLoaderTest.cpp

299 lines
12 KiB
C++
Raw Normal View History

2019-12-27 22:16:57 +08:00
//
// dataLoaderTest.cpp
// MNN
//
// Created by MNN on 2019/11/20.
// Copyright © 2018, Alibaba Group Holding Limited
//
#include <MNN/expr/ExprCreator.hpp>
#include <cmath>
#include <iostream>
#include <vector>
#include <algorithm>
#include "DataLoader.hpp"
#include "DataLoaderConfig.hpp"
#include "DemoUnit.hpp"
#include "LambdaTransform.hpp"
#include "MnistDataset.hpp"
#include "RandomSampler.hpp"
#include "StackTransform.hpp"
#include "TransformDataset.hpp"
using namespace std;
class DataLoaderTest : public DemoUnit {
public:
// this function is an example to use the lambda transform
// here we use lambda transform to normalize data from 0~255 to 0~1
static Example func(Example example) {
// an easier way to do this
auto cast = _Cast(example.data[0], halide_type_of<float>());
return {{_Multiply(cast, _Const(1.0f / 255.0f)), example.data[1]}, {example.target}};
}
virtual int run(int argc, const char* argv[]) override {
if (argc != 2) {
cout << "usage: ./runTrainDemo.out DataLoaderTest /path/to/unzipped/mnist/data/" << endl;
return 0;
}
const int testCount = 6;
int passedTestCount = 0;
std::string root = argv[1];
// train data loader
const size_t trainDatasetSize = 60000;
auto trainDataset = std::make_shared<MnistDataset>(root, MnistDataset::Mode::TRAIN);
auto trainSampler = std::make_shared<RandomSampler>(trainDataset->size());
const size_t trainBatchSize = 7;
const size_t trainNumWorkers = 4;
auto trainConfig = std::make_shared<DataLoaderConfig>(trainBatchSize, trainNumWorkers);
DataLoader trainDataLoader(trainDataset, trainSampler, trainConfig);
auto images = trainDataset->images();
auto labels = trainDataset->labels();
const int32_t kImageRows = 28;
const int32_t kImageColumns = 28;
const size_t iterations = trainDatasetSize / trainBatchSize;
auto samplerIndices = trainSampler->indices();
sort(samplerIndices.begin(), samplerIndices.end());
for (int i = 0; i < samplerIndices.size(); i++) {
MNN_ASSERT(samplerIndices[i] == i);
}
for (int i = 0; i < iterations; i++) {
auto trainData = trainDataLoader.next();
for (int j = 0; j < trainData.size(); j++) {
auto index = int(trainData[j].data[1]->readMap<float>()[0]);
auto data = trainData[j].data[0]->readMap<uint8_t>();
auto label = trainData[j].target[0]->readMap<uint8_t>();
auto trueData = images->readMap<uint8_t>() + kImageRows * kImageColumns * index;
auto trueLabel = labels->readMap<uint8_t>() + index;
for (int k = 0; k < kImageRows * kImageColumns; k++) {
MNN_ASSERT(data[k] == trueData[k]);
}
MNN_ASSERT(label[0] == trueLabel[0]);
}
}
trainDataLoader.clean();
passedTestCount++;
cout << "[" << passedTestCount << " / " << testCount << "] passed." << endl;
// the lambda transform for one example, we also can do it in batch
auto trainLambdaTransform = std::make_shared<LambdaTransform>(func);
auto trainLambdaTransDataset = std::make_shared<BatchTransformDataset>(trainDataset, trainLambdaTransform);
DataLoader trainLambdaDataLoader(trainLambdaTransDataset, trainSampler, trainConfig);
samplerIndices = trainSampler->indices();
sort(samplerIndices.begin(), samplerIndices.end());
for (int i = 0; i < samplerIndices.size(); i++) {
MNN_ASSERT(samplerIndices[i] == i);
}
std::vector<int> tempIndex;
for (int i = 0; i < iterations; i++) {
auto trainData = trainLambdaDataLoader.next();
for (int j = 0; j < trainData.size(); j++) {
auto index = int(trainData[j].data[1]->readMap<float>()[0]);
tempIndex.emplace_back(index);
auto data = trainData[j].data[0]->readMap<float>();
auto label = trainData[j].target[0]->readMap<uint8_t>();
auto trueData = images->readMap<uint8_t>() + kImageRows * kImageColumns * index;
auto trueLabel = labels->readMap<uint8_t>() + index;
for (int k = 0; k < kImageRows * kImageColumns; k++) {
MNN_ASSERT(fabs(data[k] - (trueData[k] / 255.0f)) < 1e-6);
}
MNN_ASSERT(label[0] == trueLabel[0]);
}
}
trainLambdaDataLoader.clean();
passedTestCount++;
cout << "[" << passedTestCount << " / " << testCount << "] passed." << endl;
// the stack transform, stack [1, 28, 28] to [n, 1, 28, 28]
auto trainStackTransform = std::make_shared<StackTransform>();
auto trainStackTransDataset = std::make_shared<BatchTransformDataset>(trainDataset, trainStackTransform);
DataLoader trainStackDataLoader(trainStackTransDataset, trainSampler, trainConfig);
samplerIndices = trainSampler->indices();
sort(samplerIndices.begin(), samplerIndices.end());
for (int i = 0; i < samplerIndices.size(); i++) {
MNN_ASSERT(samplerIndices[i] == i);
}
for (int i = 0; i < iterations; i++) {
auto trainData = trainStackDataLoader.next();
MNN_ASSERT(trainData.size() == 1);
std::vector<int> shape = {trainBatchSize, 1, 28, 28};
MNN_ASSERT(trainData[0].data[0]->getInfo()->dim == shape);
shape = {trainBatchSize};
MNN_ASSERT(trainData[0].target[0]->getInfo()->dim == shape);
auto data = trainData[0].data[0]->readMap<uint8_t>();
auto label = trainData[0].target[0]->readMap<uint8_t>();
for (int j = 0; j < trainBatchSize; j++) {
auto index = int(trainData[0].data[1]->readMap<float>()[j]);
auto trueData = images->readMap<uint8_t>() + kImageRows * kImageColumns * index;
auto trueLabel = labels->readMap<uint8_t>() + index;
for (int k = 0; k < kImageRows * kImageColumns; k++) {
int dataIndex = j * (kImageRows * kImageColumns) + k;
MNN_ASSERT(data[dataIndex] == trueData[k]);
}
MNN_ASSERT(label[j] == trueLabel[0]);
}
}
trainStackDataLoader.clean();
passedTestCount++;
cout << "[" << passedTestCount << " / " << testCount << "] passed." << endl;
// here we test Lambda + Stack
auto trainLambdaStackTransDataset =
std::make_shared<BatchTransformDataset>(trainLambdaTransDataset, trainStackTransform);
DataLoader trainLambdaStackDataLoader(trainLambdaStackTransDataset, trainSampler, trainConfig);
samplerIndices = trainSampler->indices();
sort(samplerIndices.begin(), samplerIndices.end());
for (int i = 0; i < samplerIndices.size(); i++) {
MNN_ASSERT(samplerIndices[i] == i);
}
for (int i = 0; i < iterations; i++) {
auto trainData = trainLambdaStackDataLoader.next();
MNN_ASSERT(trainData.size() == 1);
std::vector<int> shape = {trainBatchSize, 1, 28, 28};
MNN_ASSERT(trainData[0].data[0]->getInfo()->dim == shape);
shape = {trainBatchSize};
MNN_ASSERT(trainData[0].target[0]->getInfo()->dim == shape);
auto data = trainData[0].data[0]->readMap<float>();
auto label = trainData[0].target[0]->readMap<uint8_t>();
for (int j = 0; j < trainBatchSize; j++) {
auto index = int(trainData[0].data[1]->readMap<float>()[j]);
auto trueData = images->readMap<uint8_t>() + kImageRows * kImageColumns * index;
auto trueLabel = labels->readMap<uint8_t>() + index;
for (int k = 0; k < kImageRows * kImageColumns; k++) {
int dataIndex = j * (kImageRows * kImageColumns) + k;
MNN_ASSERT(fabs(data[dataIndex] - (trueData[k] / 255.0f)) < 1e-6);
}
MNN_ASSERT(label[j] == trueLabel[0]);
}
}
trainLambdaStackDataLoader.clean();
passedTestCount++;
cout << "[" << passedTestCount << " / " << testCount << "] passed." << endl;
// here we test Stack + Lambda
auto trainStackLambdaTransDataset =
std::make_shared<BatchTransformDataset>(trainStackTransDataset, trainLambdaTransform);
DataLoader trainStackLamdaDataLoader(trainStackLambdaTransDataset, trainSampler, trainConfig);
samplerIndices = trainSampler->indices();
sort(samplerIndices.begin(), samplerIndices.end());
for (int i = 0; i < samplerIndices.size(); i++) {
MNN_ASSERT(samplerIndices[i] == i);
}
for (int i = 0; i < iterations; i++) {
auto trainData = trainStackLamdaDataLoader.next();
MNN_ASSERT(trainData.size() == 1);
std::vector<int> shape = {trainBatchSize, 1, 28, 28};
MNN_ASSERT(trainData[0].data[0]->getInfo()->dim == shape);
shape = {trainBatchSize};
MNN_ASSERT(trainData[0].target[0]->getInfo()->dim == shape);
auto data = trainData[0].data[0]->readMap<float>();
auto label = trainData[0].target[0]->readMap<uint8_t>();
for (int j = 0; j < trainBatchSize; j++) {
auto index = int(trainData[0].data[1]->readMap<float>()[j]);
auto trueData = images->readMap<uint8_t>() + kImageRows * kImageColumns * index;
auto trueLabel = labels->readMap<uint8_t>() + index;
for (int k = 0; k < kImageRows * kImageColumns; k++) {
int dataIndex = j * (kImageRows * kImageColumns) + k;
MNN_ASSERT(fabs(data[dataIndex] - (trueData[k] / 255.0f)) < 1e-6);
}
MNN_ASSERT(label[j] == trueLabel[0]);
}
}
trainStackLamdaDataLoader.clean();
passedTestCount++;
cout << "[" << passedTestCount << " / " << testCount << "] passed." << endl;
// test makeDataLoader
auto madeDataLoader = DataLoader::makeDataLoader(
trainDataset, {nullptr, trainStackTransform, nullptr, trainLambdaTransform, nullptr}, 7);
for (int i = 0; i < iterations; i++) {
auto trainData = madeDataLoader->next();
MNN_ASSERT(trainData.size() == 1);
std::vector<int> shape = {trainBatchSize, 1, 28, 28};
MNN_ASSERT(trainData[0].data[0]->getInfo()->dim == shape);
shape = {trainBatchSize};
MNN_ASSERT(trainData[0].target[0]->getInfo()->dim == shape);
auto data = trainData[0].data[0]->readMap<float>();
auto label = trainData[0].target[0]->readMap<uint8_t>();
for (int j = 0; j < trainBatchSize; j++) {
auto index = int(trainData[0].data[1]->readMap<float>()[j]);
auto trueData = images->readMap<uint8_t>() + kImageRows * kImageColumns * index;
auto trueLabel = labels->readMap<uint8_t>() + index;
for (int k = 0; k < kImageRows * kImageColumns; k++) {
int dataIndex = j * (kImageRows * kImageColumns) + k;
MNN_ASSERT(fabs(data[dataIndex] - (trueData[k] / 255.0f)) < 1e-6);
}
MNN_ASSERT(label[j] == trueLabel[0]);
}
}
madeDataLoader->clean();
passedTestCount++;
cout << "[" << passedTestCount << " / " << testCount << "] passed." << endl;
return 0;
}
};
DemoUnitSetRegister(DataLoaderTest, "DataLoaderTest");