This module is used to read the data stored on the hard disk, package it and output it as the data type usable in MNN training. The source code of this module is located in the MNN_root/tools/train/source/data/directory. To use it, include the DataLoader.hpp header file. All other components in the module are imported to build DataLoader.

Related demo

1. MNN_root/tools/train/source/demo/dataLoaderDemo.cpp

Use the MNIST dataset to build a DataLoader and output the DataLoader.

1. MNN_root/tools/train/source/demo/dataLoaderTest.cpp

Use MNIST datasets to build DataLoader and test some components in DataLoader.

1. MNN_root/tools/train/source/demo/``ImageDatasetDemo.cpp

Reads the image data saved on the hard disk and displays it. The display needs to use OpenCV and turn it on at compile time with the ‘MNN_USE_OPENCV ‘macro.

Custom Dataset

Reference The writing method of the preset Dataset in MNN_root/tools/train/source/datasets/ inherits the Dataset class and implements two abstract functions, for example:

//  MnistDataset.cpp
// Returns an image and its label from the MNIST dataset.
Example MnistDataset::get(size_t index) {
    auto data  = _Input({1, kImageRows, kImageColumns}, NCHW, halide_type_of<uint8_t>());
    auto label = _Input({}, NCHW, halide_type_of<uint8_t>());
    auto dataPtr = mImagePtr + index * kImageRows * kImageColumns;
    ::memcpy(data->writeMap<uint8_t>(), dataPtr, kImageRows * kImageColumns);
    auto labelPtr = mLabelsPtr + index;
    ::memcpy(label->writeMap<uint8_t>(), labelPtr, 1);
    auto returnIndex = _Const(index);
    // return the index for test
    return {{data, returnIndex}, {label}};
}
// Returns the dataset size. 60000 for the training dataset, 10000 for the testing dataset.
size_t MnistDataset::size() {
    return mImages->getInfo()->dim[0];
}

DataLoader usage example

Workflow: Customize Dataset, construct DataLoader, read data, DataLoader->reset();

//
//  ImageDatasetDemo.cpp
//  MNN
//
//  Created by MNN on 2019/11/20.
//  Copyright © 2018, Alibaba Group Holding Limited
//
#include <iostream>
#include "DataLoader.hpp"
#include "DemoUnit.hpp"
#include "ImageDataset.hpp"
#include "RandomSampler.hpp"
#include "Sampler.hpp"
#include "Transform.hpp"
#include "TransformDataset.hpp"
#ifdef MNN_USE_OPENCV
#include <opencv2/opencv.hpp> // use opencv to show pictures
using namespace cv;
#endif
using namespace std;
/*
 * this is an demo for how to use the ImageDataset and DataLoader
 */
class ImageDatasetDemo : 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.first[0], halide_type_of<float>());
        example.first[0] = _Multiply(cast, _Const(1.0f / 255.0f));
        return example;
    }
    virtual int run(int argc, const char* argv[]) override {
        if (argc != 3) {
            cout << "usage: ./runTrainDemo.out ImageDatasetDemo path/to/images/ path/to/image/txt\n" << endl;
            // The data format of the ImageDataset uses the data format of ImageNet. 
            // You can define your own data set and its own data format.
            cout << "the ImageDataset read stored images as input data.\n"
                    "use 'pathToImages' and a txt file to construct a ImageDataset.\n"
                    "the txt file should use format as below:\n"
                    "     image1.jpg label1,label2,...\n"
                    "     image2.jpg label3,label4,...\n"
                    "     ...\n"
                    "the ImageDataset would read images from:\n"
                    "     pathToImages/image1.jpg\n"
                    "     pathToImages/image2.jpg\n"
                    "     ...\n"
                 << endl;
            return 0;
        }
        std::string pathToImages   = argv[1];
        std::string pathToImageTxt = argv[2];
        // Preprocess of the configurable data in ImageDataset.
        auto converImagesToFormat  = ImageDataset::DestImageFormat::RGB;
        int resizeHeight           = 224;
        int resizeWidth            = 224;
        std::vector<float> scales = {1/255.0, 1/255.0, 1/255.0};
        auto config                = ImageDataset::ImageConfig(converImagesToFormat, resizeHeight, resizeWidth, scales);
        bool readAllImagesToMemory = false;
        // Construct the ImageDataset
        auto dataset = std::make_shared<ImageDataset>(pathToImages, pathToImageTxt, config, readAllImagesToMemory);
        const int batchSize  = 1;
        const int numWorkers = 1;
        // Construct the DataLoader，
        // Here we stack the batch data into a VARP(Tensor)
        auto dataLoader = std::shared_ptr<DataLoader>(DataLoader::makeDataLoader(dataset, batchSize, true, false, numWorkers));
        const size_t iterations = dataset->size() / batchSize;
        for (int i = 0; i < iterations; i++) {
            // Reads the data.
            auto trainData = dataLoader->next();
            auto data  = trainData[0].first[0]->readMap<float_t>();
            auto label = trainData[0].second[0]->readMap<int32_t>();
            cout << "index: " << i << " label: " << int(label[0]) << endl;
#ifdef MNN_USE_OPENCV
            // only show the first picture in the batch
            Mat image = Mat(resizeHeight, resizeWidth, CV_32FC(3), (void*)data);
            imshow("image", image);
            waitKey(-1);
#endif
        }
        // DataLoader must be reset for each completed pass over the dataset.
        // This will reset the sampler's internal state
        dataLoader->reset();
        return 0;
    }
};
DemoUnitSetRegister(ImageDatasetDemo, "ImageDatasetDemo");

Related classes and concepts

VARP

Variables in MNN dynamic graphs, similar to Tensor in PyTorch.

Example

Minimum unit of data output by DataLoader

/**
 First: data: a vector of input tensors (for single input dataset is only one)
 Second: target: a vector of output tensors (for single output dataset is only one)
 */
typedef std::pair<std::vector<VARP>, std::vector<VARP>> Example;

As you can see, an Example is a data pair. The first part is input and the second part is target. Since the network may have multiple inputs and multiple targets, the first and second are vector structures.

RandomSampler : public Sampler

Random sampling sequence generator. For example, if there are 1000 images in the image dataset, the sampling sequence 0~999 is generated, and whether to shuffle or not is specified according to the configuration.

public:
    // size: length of the sampling sequence.
    // shuffle: true if the sampling sequence is to be randomly generated. 
    explicit RandomSampler(size_t size, bool shuffle = true);
    // Resets the internal state of the sampler.
    void reset(size_t size) override;
    // Size of the sampler.
    size_t size() override;
    // Returns the sampling sequence generated internally.
    const std::vector<size_t> indices();
    // Returns the number of sampling sequences already used.
    size_t index();
    // Obtains the next batchSize sampling sequence.
    std::vector<size_t> next(size_t batchSize) override;
private:
    std::vector<size_t> mIndices;
    size_t mIndex = 0;
    bool mShuffle;

Dataset

The abstract base class of a dataset. To customize a dataset, you must inherit this base class and implement abstract functions. For more information, see Description of preset datasets in MNN_root/tools/train/source/datasets/

// Returns the size of te dataset. For a dataset with 1000 images, the size is 1000.
virtual size_t size() = 0;
// Returns the data at the specified index. If `index` is 123, then it returns the 123rd image data.
virtual Example get(size_t index) = 0;
// Returns the batch of data specified by the listed of indices. 
std::vector<Example> getBatch(std::vector<size_t> indices);

Transform

Abstract base class, to perform a certain transformation on each data in a batch, can be some preprocessing, etc.

BatchTransform

Abstract base class, to perform a certain transformation on a batch of data, can be some preprocessing, etc.

StackTransform : public BatchTransform

Combine a batch of data represented by vector output from a Dataset into a VARP, that is
Stack( (c, h, w), (c, h, w), (c, h, w)... ) --> (n, c, h, w)
``

LambdaTransform : public Transform

Process each Example of the Dataset output separately, such as centralization and normalization.

TransformDataset : public Dataset

To Transform a Dataset is still a Dataset used to output data.

DataLoaderConfig

Configure DataLoader with the following configuration items:

batchSize: specifies the batch size. numWorkers: the number of threads pre-read by multiple threads

DataLoader

According to the sampling sequence generated by the sampler, obtain the corresponding data from the corresponding Dataset and output

// Constructor
DataLoader(std::shared_ptr<BatchDataset> dataset, std::shared_ptr<Sampler> sampler,
               std::shared_ptr<DataLoaderConfig> config);
// Contruct a DataLoader without Transform
static DataLoader* makeDataLoader(std::shared_ptr<BatchDataset> dataset,
                                  const int batchSize,
                                  const bool stack = true, // Whether to stack a batch of data into a VARP(TENSOR)
                                  const bool shuffle = true,
                                  const int numWorkers = 0);
// Construct a DataLoader with Transform，multiple Transforms can be stacked.
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);
// The number of iterations after specifying the batch size. The latest batch can be smaller than the 
// batchsize. 
size_t iterNumber() const;
// The size of the dataset.
size_t size() const;
// Gets the data for the next batch.
std::vector<Example> next();
// Clear internal data array, and resets the internal sampler.
void clean();
// clean() + resets the dataset. Every time the dataset completes a full round of data
// output, it must be reset.
void reset();