io.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
 
#ifndef MXNET_IO_H_
#define MXNET_IO_H_
 
#include <vector>
#include <string>
#include <utility>
#include <queue>
#include "dmlc/data.h"
#include "dmlc/registry.h"
#include "./base.h"
#include "./ndarray.h"
 
namespace mxnet {
template <typename DType>
class IIterator : public dmlc::DataIter<DType> {
 public:
  virtual void Init(const std::vector<std::pair<std::string, std::string> >& kwargs) = 0;
  virtual void BeforeFirst(void) = 0;
  virtual bool Next(void) = 0;
  virtual const DType& Value(void) const = 0;
  virtual ~IIterator(void) {}
  std::vector<std::string> data_names;
  inline void SetDataName(const std::string data_name) {
    data_names.push_back(data_name);
  }
  virtual int64_t GetLenHint(void) const {
    return -1;
  }
};  // class IIterator
 
struct DataInst {
  unsigned index;
  std::vector<TBlob> data;
  std::string extra_data;
};  // struct DataInst
 
struct DataBatch {
  std::vector<NDArray> data;
  std::vector<uint64_t> index;
  std::string extra_data;
  int num_batch_padd;
};  // struct DataBatch
 
typedef std::function<IIterator<DataBatch>*()> DataIteratorFactory;
struct DataIteratorReg : public dmlc::FunctionRegEntryBase<DataIteratorReg, DataIteratorFactory> {};
//--------------------------------------------------------------
// The following part are API Registration of Iterators
//--------------------------------------------------------------
#define MXNET_REGISTER_IO_ITER(name) \
  DMLC_REGISTRY_REGISTER(::mxnet::DataIteratorReg, DataIteratorReg, name)
 
class Dataset {
 public:
  virtual uint64_t GetLen(void) const = 0;
  virtual bool GetItem(uint64_t idx, std::vector<NDArray>* ret) = 0;
  // virtual destructor
  virtual ~Dataset(void) {}
};  // class Dataset
 
typedef std::function<Dataset*(const std::vector<std::pair<std::string, std::string> >&)>
    DatasetFactory;
struct DatasetReg : public dmlc::FunctionRegEntryBase<DatasetReg, DatasetFactory> {};
//--------------------------------------------------------------
// The following part are API Registration of Datasets
//--------------------------------------------------------------
#define MXNET_REGISTER_IO_DATASET(name) \
  DMLC_REGISTRY_REGISTER(::mxnet::DatasetReg, DatasetReg, name)
 
class BatchifyFunction {
 public:
  virtual ~BatchifyFunction(void) {}
  virtual bool Batchify(const std::vector<std::vector<NDArray> >& inputs,
                        std::vector<NDArray>* outputs) = 0;
};  // class BatchifyFunction
 
using BatchifyFunctionPtr = std::shared_ptr<BatchifyFunction>;
 
typedef std::function<BatchifyFunction*(const std::vector<std::pair<std::string, std::string> >&)>
    BatchifyFunctionFactory;
struct BatchifyFunctionReg
    : public dmlc::FunctionRegEntryBase<BatchifyFunctionReg, BatchifyFunctionFactory> {};
//--------------------------------------------------------------
// The following part are API Registration of Batchify Function
//--------------------------------------------------------------
#define MXNET_REGISTER_IO_BATCHIFY_FUNCTION(name) \
  DMLC_REGISTRY_REGISTER(::mxnet::BatchifyFunctionReg, BatchifyFunctionReg, name)
}  // namespace mxnet
#endif  // MXNET_IO_H_