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);
  }
};  // 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)
}  // namespace mxnet
#endif  // MXNET_IO_H_