You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2018/01/31 19:30:33 UTC
[GitHub] piiswrong closed pull request #8437: Optimize the performance of mxnet.image.ImageIter
piiswrong closed pull request #8437: Optimize the performance of mxnet.image.ImageIter
URL: https://github.com/apache/incubator-mxnet/pull/8437
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/Makefile b/Makefile
index e821c6faa5..ecd7f5ea59 100644
--- a/Makefile
+++ b/Makefile
@@ -318,6 +318,28 @@ else
endif
endif
+ifeq ($(USE_PB_FORMAT_REC), 1)
+ifeq ($(USE_DIST_KVSTORE), 0)
+$(error "USE_DIST_KVSTORE must be setted to 1 when USE_PB_FORMAT_REC is 1")
+endif
+ifeq ($(USE_OPENMP), 1)
+$(error "USE_OPENMP must be setted to 0 when USE_PB_FORMAT_REC is 1")
+endif
+
+ifndef PROTOC
+PROTOC = ${DEPS_PATH}/bin/protoc
+endif
+
+CFLAGS += -DPB_FORMAT_REC
+OBJ += build/src/io/recordio.pb.o
+SRC_TARGET = src/io/recordio.pb.h
+
+include $(PS_PATH)/make/deps.mk
+src/io/recordio.pb.cc src/io/recordio.pb.h : src/io/recordio.proto $(PROTOBUF)
+ $(PROTOC) --cpp_out=./src/io/ --proto_path=./src/io/ $<
+ $(PROTOC) --python_out=./tools/ --proto_path=./src/io/ $<
+endif
+
# all dep
LIB_DEP += $(DMLC_CORE)/libdmlc.a $(NNVM_PATH)/lib/libnnvm.a
ALL_DEP = $(OBJ) $(EXTRA_OBJ) $(PLUGIN_OBJ) $(LIB_DEP)
@@ -345,7 +367,7 @@ endif
# For quick compile test, used smaller subset
ALLX_DEP= $(ALL_DEP)
-build/src/%.o: src/%.cc
+build/src/%.o: src/%.cc $(SRC_TARGET)
@mkdir -p $(@D)
$(CXX) -std=c++11 -c $(CFLAGS) -MMD -c $< -o $@
@@ -512,6 +534,7 @@ ifneq ($(EXTRA_OPERATORS),)
clean: cyclean $(EXTRA_PACKAGES_CLEAN)
$(RM) -r build lib bin *~ */*~ */*/*~ */*/*/*~ R-package/NAMESPACE R-package/man R-package/R/mxnet_generated.R \
R-package/inst R-package/src/*.o R-package/src/*.so mxnet_*.tar.gz
+ find src -name "*.pb.[ch]*" -delete
cd $(DMLC_CORE); $(MAKE) clean; cd -
cd $(PS_PATH); $(MAKE) clean; cd -
cd $(NNVM_PATH); $(MAKE) clean; cd -
@@ -522,6 +545,7 @@ else
clean: cyclean testclean $(EXTRA_PACKAGES_CLEAN)
$(RM) -r build lib bin *~ */*~ */*/*~ */*/*/*~ R-package/NAMESPACE R-package/man R-package/R/mxnet_generated.R \
R-package/inst R-package/src/image_recordio.h R-package/src/*.o R-package/src/*.so mxnet_*.tar.gz
+ find src -name "*.pb.[ch]*" -delete
cd $(DMLC_CORE); $(MAKE) clean; cd -
cd $(PS_PATH); $(MAKE) clean; cd -
cd $(NNVM_PATH); $(MAKE) clean; cd -
diff --git a/dmlc-core b/dmlc-core
index 595d02c0e8..2f4b61ac06 160000
--- a/dmlc-core
+++ b/dmlc-core
@@ -1 +1 @@
-Subproject commit 595d02c0e87be8a0846700462b6f45f1b1031e39
+Subproject commit 2f4b61ac06087254ab492e10f9a2e263ba83112b
diff --git a/example/image-classification/common/data.py b/example/image-classification/common/data.py
index dc8915cda4..03031091c0 100755
--- a/example/image-classification/common/data.py
+++ b/example/image-classification/common/data.py
@@ -19,6 +19,7 @@
import random
from mxnet.io import DataBatch, DataIter
import numpy as np
+import os
def add_data_args(parser):
data = parser.add_argument_group('Data', 'the input images')
@@ -38,6 +39,10 @@ def add_data_args(parser):
help='number of threads for data decoding')
data.add_argument('--benchmark', type=int, default=0,
help='if 1, then feed the network with synthetic data')
+ data.add_argument('--use-python-iter', type=int, default=0,
+ help='whether or not use python version iter')
+ data.add_argument('--use-gpu-augmenter', type=int, default=0,
+ help='whether or not augment the image in gpus')
return data
def add_data_aug_args(parser):
@@ -118,50 +123,90 @@ def get_rec_iter(args, kv=None):
(rank, nworker) = (kv.rank, kv.num_workers)
else:
(rank, nworker) = (0, 1)
- rgb_mean = [float(i) for i in args.rgb_mean.split(',')]
- train = mx.io.ImageRecordIter(
- path_imgrec = args.data_train,
- path_imgidx = args.data_train_idx,
- label_width = 1,
- mean_r = rgb_mean[0],
- mean_g = rgb_mean[1],
- mean_b = rgb_mean[2],
- data_name = 'data',
- label_name = 'softmax_label',
- data_shape = image_shape,
- batch_size = args.batch_size,
- rand_crop = args.random_crop,
- max_random_scale = args.max_random_scale,
- pad = args.pad_size,
- fill_value = 127,
- min_random_scale = args.min_random_scale,
- max_aspect_ratio = args.max_random_aspect_ratio,
- random_h = args.max_random_h,
- random_s = args.max_random_s,
- random_l = args.max_random_l,
- max_rotate_angle = args.max_random_rotate_angle,
- max_shear_ratio = args.max_random_shear_ratio,
- rand_mirror = args.random_mirror,
- preprocess_threads = args.data_nthreads,
- shuffle = True,
- num_parts = nworker,
- part_index = rank)
+ rgb_mean = np.array([float(i) for i in args.rgb_mean.split(',')])
+ gpu_devs = None
+ if args.use_gpu_augmenter:
+ gpu_devs = None if args.gpus is None or args.gpus is '' else [
+ mx.gpu(int(i)) for i in args.gpus.split(',')]
+ if not args.use_python_iter:
+ train = mx.io.ImageRecordIter(
+ path_imgrec = args.data_train,
+ path_imgidx = args.data_train_idx,
+ label_width = 1,
+ mean_r = rgb_mean[0],
+ mean_g = rgb_mean[1],
+ mean_b = rgb_mean[2],
+ data_name = 'data',
+ label_name = 'softmax_label',
+ data_shape = image_shape,
+ batch_size = args.batch_size,
+ rand_crop = args.random_crop,
+ max_random_scale = args.max_random_scale,
+ pad = args.pad_size,
+ fill_value = 127,
+ min_random_scale = args.min_random_scale,
+ max_aspect_ratio = args.max_random_aspect_ratio,
+ random_h = args.max_random_h,
+ random_s = args.max_random_s,
+ random_l = args.max_random_l,
+ max_rotate_angle = args.max_random_rotate_angle,
+ max_shear_ratio = args.max_random_shear_ratio,
+ rand_mirror = args.random_mirror,
+ preprocess_threads = args.data_nthreads,
+ shuffle = True,
+ num_parts = nworker,
+ part_index = rank)
+ else:
+ train = mx.img.PrefetchImageIter(
+ path_imgrec = args.data_train,
+ pb_format = True,
+ mean = rgb_mean,
+ data_name = 'data',
+ label_name = 'softmax_label',
+ data_shape = image_shape,
+ batch_size = args.batch_size,
+ rand_crop = args.random_crop,
+ rand_mirror = args.random_mirror,
+ std = True,
+ preprocess_threads = args.data_nthreads,
+ shuffle = True,
+ num_parts = nworker,
+ part_index = rank,
+ gpu_ctxs = gpu_devs)
+
if args.data_val is None:
return (train, None)
- val = mx.io.ImageRecordIter(
- path_imgrec = args.data_val,
- path_imgidx = args.data_val_idx,
- label_width = 1,
- mean_r = rgb_mean[0],
- mean_g = rgb_mean[1],
- mean_b = rgb_mean[2],
- data_name = 'data',
- label_name = 'softmax_label',
- batch_size = args.batch_size,
- data_shape = image_shape,
- preprocess_threads = args.data_nthreads,
- rand_crop = False,
- rand_mirror = False,
- num_parts = nworker,
- part_index = rank)
+ if not args.use_python_iter:
+ val = mx.io.ImageRecordIter(
+ path_imgrec = args.data_val,
+ path_imgidx = args.data_val_idx,
+ label_width = 1,
+ mean_r = rgb_mean[0],
+ mean_g = rgb_mean[1],
+ mean_b = rgb_mean[2],
+ data_name = 'data',
+ label_name = 'softmax_label',
+ batch_size = args.batch_size,
+ data_shape = image_shape,
+ preprocess_threads = args.data_nthreads,
+ rand_crop = False,
+ rand_mirror = False,
+ round_batch = False,
+ num_parts = nworker,
+ part_index = rank)
+ else:
+ val = mx.img.PrefetchImageIter(
+ path_imgrec = args.data_val,
+ pb_format = True,
+ mean = rgb_mean,
+ data_name = 'data',
+ label_name = 'softmax_label',
+ data_shape = image_shape,
+ batch_size = args.batch_size,
+ rand_crop = False,
+ rand_mirror = False,
+ std = True,
+ num_parts = nworker,
+ part_index = rank,
+ gpu_ctxs = gpu_devs)
return (train, val)
diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h
index 55b840dd2c..841c271c57 100644
--- a/include/mxnet/c_api.h
+++ b/include/mxnet/c_api.h
@@ -88,6 +88,8 @@ typedef void *RtcHandle;
typedef void *CudaModuleHandle;
/*! \brief handle to rtc cuda kernel*/
typedef void *CudaKernelHandle;
+/*! \brief handle to RecordIter*/
+typedef void *RecordIterHandle;
typedef void (*ExecutorMonitorCallback)(const char*,
NDArrayHandle,
@@ -1995,6 +1997,150 @@ MXNET_DLL int MXRtcCudaKernelCall(CudaKernelHandle handle, int dev_id, void** ar
mx_uint block_dim_y, mx_uint block_dim_z,
mx_uint shared_mem);
+
+/*!
+ * \brief Get the number of data and label blobs. By Alan
+ * \param handle the handle pointer to the data iterator
+ * \param data_size the number of data blobs
+ * \param label_size the number of label blobs
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXDataIterGetDataAndLabelSize(DataIterHandle handle,
+ mx_uint *data_size,
+ mx_uint* label_size);
+
+/*!
+ * \brief Get the detailed information about record iterator.
+ * \param name The returned name of the record iter.
+ * \param description The returned description of the iter.
+ * \param num_args Number of arguments.
+ * \param arg_names Name of the arguments.
+ * \param arg_type_infos Type informations about the arguments.
+ * \param arg_descriptions Description information about the arguments.
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXRecordIterGetIterInfo(const char **name,
+ const char **description,
+ mx_uint *num_args,
+ const char ***arg_names,
+ const char ***arg_type_infos,
+ const char ***arg_descriptions);
+
+/*!
+ * \brief Init an iterator, init with parameters
+ * the array size of passed in arguments
+ * \param num_param number of parameter
+ * \param keys parameter keys
+ * \param vals parameter values
+ * \param out resulting iterator
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXRecordIterCreate(mx_uint num_param,
+ const char **keys,
+ const char **vals,
+ RecordIterHandle *out);
+/*!
+ * \brief Free the handle to the IO module
+ * \param handle the handle pointer to the record iterator
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXRecordIterFree(RecordIterHandle handle);
+
+/*!
+ * \brief Call iterator.Reset
+ * \param handle the handle to iterator
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXRecordIterBeforeFirst(RecordIterHandle handle);
+
+/*!
+ * \brief Move iterator to next position
+ * \param handle the handle to iterator
+ * \param out return value of next
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXRecordIterNext(RecordIterHandle handle, int *out);
+
+/*!
+ * \brief Get num of data
+ * \param handle the handle of the record iterator
+ * \param out number of data.
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXRecordIterGetDataNum(RecordIterHandle handle,
+ mx_uint *out);
+
+/*!
+ * \brief Get record data
+ * \param handle the handle of the record iterator
+ * \param index the item index in data
+ * \param out_id the id of item
+ * \param out_type item type
+ * \param out_ndarray NDArray value of item
+ * \param out_string_len len of string value
+ * \param out_string string value
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXRecordIterGetData(RecordIterHandle handle,
+ mx_uint index,
+ uint64_t *out_id,
+ int *out_type,
+ NDArrayHandle *out_ndarray,
+ mx_uint *out_string_len,
+ const char **out_string);
+
+
+/*!
+ * \brief Get num of extra data
+ * \param handle the handle of the record iterator
+ * \param out number of extra data.
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXRecordIterGetExtraDataNum(RecordIterHandle handle,
+ mx_uint *out);
+
+/*!
+ * \brief Get record extra data
+ * \param handle the handle of the record iterator
+ * \param index the item index in extra
+ * \param out_key item key
+ * \param out_key_len len of item key
+ * \param out_type item type
+ * \param out_ndarray NDArray value of item
+ * \param out_string_len len of string value
+ * \param out_string string value
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXRecordIterGetExtraData(RecordIterHandle handle,
+ mx_uint index,
+ const char **out_key,
+ mx_uint *out_key_len,
+ int *out_type,
+ NDArrayHandle *out_ndarray,
+ const char **out_string,
+ mx_uint *out_string_len);
+
+/*!
+ * \brief Get list of underlying label
+ * \param handle the handle of the record iterator
+ * \param out head of underlying label mx_float list
+ * \param out_size number of label list.
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXRecordIterGetLabel(RecordIterHandle handle,
+ mx_float **out,
+ mx_uint *out_size);
+
+/*!
+ * \brief Get the record head.
+ * \param handle the handle of the record iterator
+ * \param out_id the pointer to record index
+ * \param out_reserve the pointer to record reserve
+ * \return 0 when success, -1 when failure happens
+ */
+MXNET_DLL int MXRecordIterGetHead(RecordIterHandle handle,
+ uint64_t *out_id,
+ uint64_t *out_reserve);
#ifdef __cplusplus
}
#endif // __cplusplus
diff --git a/make/config.mk b/make/config.mk
index d47d4d6931..a9294436a7 100644
--- a/make/config.mk
+++ b/make/config.mk
@@ -56,6 +56,8 @@ USE_CUDA_PATH = NONE
# whether use CuDNN R3 library
USE_CUDNN = 0
+# whether use cuSOLVER library
+USE_CUSOLVER = 0
# whether use opencv during compilation
# you can disable it, however, you will not able to use
@@ -167,6 +169,10 @@ EXTRA_OPERATORS =
# Create C++ interface package
USE_CPP_PACKAGE = 0
+# Running with the Protobuf-format rec files. If yes, then USE_DIST_KVSTORE must be setted to 1 and
+# USE_OPENMP be setted to 0
+USE_PB_FORMAT_REC = 0
+
#----------------------------
# plugins
#----------------------------
diff --git a/mshadow b/mshadow
index cb5c9872b5..87e4e1be5b 160000
--- a/mshadow
+++ b/mshadow
@@ -1 +1 @@
-Subproject commit cb5c9872b542220be5b99f3aca0e1ff56e31b028
+Subproject commit 87e4e1be5bf1875640ee9b062c8760152689b54e
diff --git a/ps-lite b/ps-lite
index bdd4c67e9e..01a81d3c0b 160000
--- a/ps-lite
+++ b/ps-lite
@@ -1 +1 @@
-Subproject commit bdd4c67e9e34dc0b8350ce306b0caa737eb31c83
+Subproject commit 01a81d3c0b15e85572444d7fa7702f39d935d8bc
diff --git a/python/mxnet/base.py b/python/mxnet/base.py
index 138bb4d157..0d570e886a 100644
--- a/python/mxnet/base.py
+++ b/python/mxnet/base.py
@@ -123,6 +123,7 @@ def _load_lib():
ExecutorHandle = ctypes.c_void_p
DataIterCreatorHandle = ctypes.c_void_p
DataIterHandle = ctypes.c_void_p
+RecordIterHandle = ctypes.c_void_p
KVStoreHandle = ctypes.c_void_p
RecordIOHandle = ctypes.c_void_p
RtcHandle = ctypes.c_void_p
diff --git a/python/mxnet/image/image.py b/python/mxnet/image/image.py
index ace9cb18ff..036d0a7652 100644
--- a/python/mxnet/image/image.py
+++ b/python/mxnet/image/image.py
@@ -26,6 +26,8 @@
import logging
import json
import numpy as np
+import Queue
+import threading
try:
import cv2
@@ -39,6 +41,7 @@
from ..ndarray._internal import _cvcopyMakeBorder as copyMakeBorder
from .. import io
from .. import recordio
+from ..context import cpu
def imread(filename, *args, **kwargs):
@@ -478,6 +481,21 @@ def random_size_crop(src, size, min_area, ratio, interp=2):
# fall back to center_crop
return center_crop(src, size, interp)
+class MultiContextConstArray(object):
+ """multi-context array.
+
+ This array will buffer the copy in multi-contexts.
+ """
+ def __init__(self, arr):
+ self.ori_arr = arr
+ self.ctx_dict = {arr.context: arr}
+
+ def get(self, ctx):
+ ret_arr = self.ctx_dict.get(ctx)
+ if ret_arr is None:
+ ret_arr = self.ori_arr.as_in_context(ctx)
+ self.ctx_dict[ctx] = ret_arr
+ return ret_arr
class Augmenter(object):
"""Image Augmenter base class"""
@@ -690,12 +708,12 @@ class ContrastJitterAug(Augmenter):
def __init__(self, contrast):
super(ContrastJitterAug, self).__init__(contrast=contrast)
self.contrast = contrast
- self.coef = nd.array([[[0.299, 0.587, 0.114]]])
+ self.coef = MultiContextConstArray(nd.array([[[0.299, 0.587, 0.114]]]))
def __call__(self, src):
"""Augmenter body"""
alpha = 1.0 + random.uniform(-self.contrast, self.contrast)
- gray = src * self.coef
+ gray = src * self.coef.get(ctx=src.context)
gray = (3.0 * (1.0 - alpha) / gray.size) * nd.sum(gray)
src *= alpha
src += gray
@@ -713,12 +731,12 @@ class SaturationJitterAug(Augmenter):
def __init__(self, saturation):
super(SaturationJitterAug, self).__init__(saturation=saturation)
self.saturation = saturation
- self.coef = nd.array([[[0.299, 0.587, 0.114]]])
+ self.coef = MultiContextConstArray(nd.array([[[0.299, 0.587, 0.114]]]))
def __call__(self, src):
"""Augmenter body"""
alpha = 1.0 + random.uniform(-self.saturation, self.saturation)
- gray = src * self.coef
+ gray = src * self.coef.get(ctx=src.context)
gray = nd.sum(gray, axis=2, keepdims=True)
gray *= (1.0 - alpha)
src *= alpha
@@ -756,7 +774,7 @@ def __call__(self, src):
[0.0, u, -w],
[0.0, w, u]])
t = np.dot(np.dot(self.ityiq, bt), self.tyiq).T
- src = nd.dot(src, nd.array(t))
+ src = nd.dot(src, nd.array(t, ctx=src.context))
return src
@@ -805,7 +823,7 @@ def __call__(self, src):
"""Augmenter body"""
alpha = np.random.normal(0, self.alphastd, size=(3,))
rgb = np.dot(self.eigvec * alpha, self.eigval)
- src += nd.array(rgb)
+ src += nd.array(rgb, ctx=src.context)
return src
@@ -821,12 +839,14 @@ class ColorNormalizeAug(Augmenter):
"""
def __init__(self, mean, std):
super(ColorNormalizeAug, self).__init__(mean=mean, std=std)
- self.mean = nd.array(mean) if mean is not None else None
- self.std = nd.array(std) if std is not None else None
+ self.mean = MultiContextConstArray(nd.array(mean)) if mean is not None else None
+ self.std = MultiContextConstArray(nd.array(std)) if std is not None else None
def __call__(self, src):
"""Augmenter body"""
- return color_normalize(src, self.mean, self.std)
+ return color_normalize(src,
+ self.mean.get(src.context) if self.mean is not None else None,
+ self.std.get(src.context) if self.std is not None else None)
class RandomGrayAug(Augmenter):
@@ -840,17 +860,16 @@ class RandomGrayAug(Augmenter):
def __init__(self, p):
super(RandomGrayAug, self).__init__(p=p)
self.p = p
- self.mat = nd.array([[0.21, 0.21, 0.21],
+ self.mat = MultiContextConstArray(nd.array([[0.21, 0.21, 0.21],
[0.72, 0.72, 0.72],
- [0.07, 0.07, 0.07]])
+ [0.07, 0.07, 0.07]]))
def __call__(self, src):
"""Augmenter body"""
if random.random() < self.p:
- src = nd.dot(src, self.mat)
+ src = nd.dot(src, self.mat.get(src.context))
return src
-
class HorizontalFlipAug(Augmenter):
"""Random horizontal flip.
@@ -881,10 +900,27 @@ def __call__(self, src):
src = src.astype(self.typ)
return src
+class CopyGpuAug(Augmenter):
+ """Copy to device to do following augs.
+
+ Parameters
+ ----------
+ ctxs : list of Context
+ GPU contexs
+ """
+ def __init__(self, ctxs):
+ super(CopyGpuAug, self).__init__(ctxs=ctxs)
+ self.ctxs = ctxs
+
+ def __call__(self, src):
+ """Augmenter body"""
+ idx = random.randint(0, len(self.ctxs)-1)
+ src = src.copyto(self.ctxs[idx])
+ return src
def CreateAugmenter(data_shape, resize=0, rand_crop=False, rand_resize=False, rand_mirror=False,
mean=None, std=None, brightness=0, contrast=0, saturation=0, hue=0,
- pca_noise=0, rand_gray=0, inter_method=2):
+ pca_noise=0, rand_gray=0, inter_method=2, gpu_ctxs=None, **kwargs):
"""Creates an augmenter list.
Parameters
@@ -934,6 +970,10 @@ def CreateAugmenter(data_shape, resize=0, rand_crop=False, rand_resize=False, ra
When shrinking an image, it will generally look best with AREA-based
interpolation, whereas, when enlarging an image, it will generally look best
with Bicubic (slow) or Bilinear (faster but still looks OK).
+ gpu_ctxs : list of Context or None
+ GPU contexts.
+ kwargs : ...
+ Other arguments.
Examples
--------
@@ -959,6 +999,9 @@ def CreateAugmenter(data_shape, resize=0, rand_crop=False, rand_resize=False, ra
else:
auglist.append(CenterCropAug(crop_size, inter_method))
+ if gpu_ctxs is not None:
+ auglist.append(CopyGpuAug(gpu_ctxs))
+
if rand_mirror:
auglist.append(HorizontalFlipAug(0.5))
@@ -1003,7 +1046,8 @@ class ImageIter(io.DataIter):
To load input images from .rec files, use `path_imgrec` parameter and to load from raw image
files, use `path_imglist` and `path_root` parameters.
- To use data partition (for distributed training) or shuffling, specify `path_imgidx` parameter.
+ To use data partition (for distributed training) or shuffling, specify `path_imgidx` parameter,
+ But when using pb-format data, `path_imgidx` is not necessary.
Parameters
----------
@@ -1028,6 +1072,8 @@ class ImageIter(io.DataIter):
Root folder of image files.
path_imgidx : str
Path to image index file. Needed for partition and shuffling when using .rec source.
+ pb_format : bool
+ whether use pb format data
shuffle : bool
Whether to shuffle all images at the start of each iteration or not.
Can be slow for HDD.
@@ -1040,30 +1086,44 @@ class ImageIter(io.DataIter):
label_name : str
Label name for provided symbols.
kwargs : ...
- More arguments for creating augmenter. See mx.image.CreateAugmenter.
+ More arguments for creating augmenter and RecordIter. See mx.image.CreateAugmenter
+ and mx.io.RecordIter.
"""
def __init__(self, batch_size, data_shape, label_width=1,
path_imgrec=None, path_imglist=None, path_root=None, path_imgidx=None,
+ pb_format=False,
shuffle=False, part_index=0, num_parts=1, aug_list=None, imglist=None,
data_name='data', label_name='softmax_label', **kwargs):
super(ImageIter, self).__init__()
assert path_imgrec or path_imglist or (isinstance(imglist, list))
+ if pb_format and kwargs.has_key('preprocess_threads'):
+ os.environ["MXNET_CPU_WORKER_NTHREADS"] = str(kwargs.pop('preprocess_threads'))
num_threads = os.environ.get('MXNET_CPU_WORKER_NTHREADS', 1)
logging.info('Using %s threads for decoding...', str(num_threads))
logging.info('Set enviroment variable MXNET_CPU_WORKER_NTHREADS to a'
' larger number to use more threads.')
class_name = self.__class__.__name__
+ self.pb_format = pb_format
if path_imgrec:
logging.info('%s: loading recordio %s...',
class_name, path_imgrec)
- if path_imgidx:
- self.imgrec = recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
- self.imgidx = list(self.imgrec.keys)
+ if self.pb_format:
+ self.recorditer = recordio.RecordIter(
+ path_imgrec = path_imgrec,
+ shuffle = shuffle,
+ num_parts = num_parts,
+ part_index = part_index,
+ **kwargs)
else:
- self.imgrec = recordio.MXRecordIO(path_imgrec, 'r') # pylint: disable=redefined-variable-type
- self.imgidx = None
+ if path_imgidx:
+ self.imgrec = recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r') # pylint: disable=redefined-variable-type
+ self.imgidx = list(self.imgrec.keys)
+ else:
+ self.imgrec = recordio.MXRecordIO(path_imgrec, 'r') # pylint: disable=redefined-variable-type
+ self.imgidx = None
else:
+ self.recorditer = None
self.imgrec = None
if path_imglist:
@@ -1109,65 +1169,86 @@ def __init__(self, batch_size, data_shape, label_width=1,
self.data_shape = data_shape
self.label_width = label_width
- self.shuffle = shuffle
- if self.imgrec is None:
- self.seq = imgkeys
- elif shuffle or num_parts > 1:
- assert self.imgidx is not None
- self.seq = self.imgidx
- else:
- self.seq = None
+ if not self.pb_format:
+ self.shuffle = shuffle
+ if self.imgrec is None:
+ self.seq = imgkeys
+ elif shuffle or num_parts > 1:
+ assert self.imgidx is not None
+ self.seq = self.imgidx
+ else:
+ self.seq = None
- if num_parts > 1:
- assert part_index < num_parts
- N = len(self.seq)
- C = N // num_parts
- self.seq = self.seq[part_index * C:(part_index + 1) * C]
+ if num_parts > 1:
+ assert part_index < num_parts
+ N = len(self.seq)
+ C = N // num_parts
+ self.seq = self.seq[part_index * C:(part_index + 1) * C]
if aug_list is None:
self.auglist = CreateAugmenter(data_shape, **kwargs)
else:
self.auglist = aug_list
- self.cur = 0
+ self.pre_batch_data = None
+ self.pre_batch_label = None
self.reset()
+ self.cur = 0
def reset(self):
"""Resets the iterator to the beginning of the data."""
- if self.shuffle:
- random.shuffle(self.seq)
- if self.imgrec is not None:
- self.imgrec.reset()
+ self.pre_batch_data = None
+ self.pre_batch_label = None
+ if self.pb_format:
+ if self.recorditer is not None:
+ self.recorditer.reset()
+ else:
+ if self.shuffle:
+ random.shuffle(self.seq)
+ if self.imgrec is not None:
+ self.imgrec.reset()
self.cur = 0
def next_sample(self):
"""Helper function for reading in next sample."""
- if self.seq is not None:
- if self.cur >= len(self.seq):
- raise StopIteration
- idx = self.seq[self.cur]
- self.cur += 1
- if self.imgrec is not None:
- s = self.imgrec.read_idx(idx)
- header, img = recordio.unpack(s)
- if self.imglist is None:
- return header.label, img
+ if self.pb_format:
+ s = self.recorditer.next()
+ assert s.data[0].type == recordio.RecordIOType.BINARY, 'this data must be binary data'
+ self.cur = s.header.id
+ return s.label[0], s.data[0].nd_value
+ else:
+ if self.seq is not None:
+ if self.cur >= len(self.seq):
+ raise StopIteration
+ idx = self.seq[self.cur]
+ self.cur += 1
+ if self.imgrec is not None:
+ s = self.imgrec.read_idx(idx)
+ header, img = recordio.unpack(s)
+ if self.imglist is None:
+ return header.label, img
+ else:
+ return self.imglist[idx][0], img
else:
- return self.imglist[idx][0], img
+ label, fname = self.imglist[idx]
+ return label, self.read_image(fname)
else:
- label, fname = self.imglist[idx]
- return label, self.read_image(fname)
- else:
- s = self.imgrec.read()
- if s is None:
- raise StopIteration
- header, img = recordio.unpack(s)
- return header.label, img
+ s = self.imgrec.read()
+ if s is None:
+ raise StopIteration
+ header, img = recordio.unpack(s)
+ return header.label, img
- def next(self):
+ def next_batch(self, data_batch=None):
"""Returns the next batch of data."""
batch_size = self.batch_size
c, h, w = self.data_shape
- batch_data = nd.empty((batch_size, c, h, w))
- batch_label = nd.empty(self.provide_label[0][1])
+ if data_batch is None:
+ batch_data = nd.empty((batch_size, c, h, w))
+ batch_label = nd.empty(self.provide_label[0][1])
+ else:
+ batch_data = data_batch.data[0]
+ batch_label = data_batch.label[0]
+ assert batch_data.shape == (batch_size, c, h, w), "batch_data shape mismatch"
+ assert batch_label.shape == self.provide_label[0][1], "batch_label shape mismatch"
i = 0
try:
while i < batch_size:
@@ -1183,12 +1264,25 @@ def next(self):
batch_data[i] = self.postprocess_data(data)
batch_label[i] = label
i += 1
+ # store the last batch_data
+ self.pre_batch_data = batch_data
+ self.pre_batch_label = batch_label
except StopIteration:
if not i:
raise StopIteration
+ else:
+ if self.pre_batch_data is None:
+ raise ValueError('must have one batch')
+ else:
+ # pad current batch with the latest batch_data
+ batch_data[i:] = self.pre_batch_data[i:]
+ batch_label[i:] = self.pre_batch_label[i:]
return io.DataBatch([batch_data], [batch_label], batch_size - i)
+ def next(self):
+ return self.next_batch()
+
def check_data_shape(self, data_shape):
"""Checks if the input data shape is valid"""
if not len(data_shape) == 3:
@@ -1206,16 +1300,19 @@ def imdecode(self, s):
See mx.img.imdecode for more details."""
def locate():
"""Locate the image file/index if decode fails."""
- if self.seq is not None:
- idx = self.seq[self.cur - 1]
- else:
- idx = self.cur - 1
- if self.imglist is not None:
- _, fname = self.imglist[idx]
- msg = "filename: {}".format(fname)
+ if self.pb_format:
+ return "Broken image index: {}".format(self.cur)
else:
- msg = "index: {}".format(idx)
- return "Broken image " + msg
+ if self.seq is not None:
+ idx = self.seq[self.cur - 1]
+ else:
+ idx = self.cur - 1
+ if self.imglist is not None:
+ _, fname = self.imglist[idx]
+ msg = "filename: {}".format(fname)
+ else:
+ msg = "index: {}".format(idx)
+ return "Broken image " + msg
try:
img = imdecode(s)
except Exception as e:
@@ -1241,4 +1338,155 @@ def augmentation_transform(self, data):
def postprocess_data(self, datum):
"""Final postprocessing step before image is loaded into the batch."""
- return nd.transpose(datum, axes=(2, 0, 1))
+ return nd.transpose(datum, axes=(2, 0, 1)).as_in_context(cpu())
+
+class PrefetchDataFlag(object):
+ ITER_COMPLETE = 0
+
+class PrefetchThreadIter(io.DataIter):
+ """Performs pre-fetch for other data iterator.
+
+ This iterator will create another thread to perform ``iter_next`` and then
+ store the data in queue. It potentially accelerates the data read, at the
+ cost of more memory usage.
+
+ Parameters
+ ----------
+ data_iter : DataIter
+ The data iterators to be pre-fetched.
+ capacity : int
+ max prefetch num
+
+ Examples
+ --------
+ >>> iter1 = mx.io.NDArrayIter({'data':mx.nd.ones((100,10))}, batch_size=25)
+ >>> piter = mx.img.PrefetchThreadIter(iter1)
+ >>> print(piter.provide_data)
+ DataDesc[data_1,(25, 10L),<type 'numpy.float32'>,NCHW]
+ """
+ def __init__(self, data_iter, capacity = 4):
+ super(PrefetchThreadIter, self).__init__()
+ assert data_iter is not None
+ self.data_iter = data_iter
+ self.batch_size = self.provide_data[0][1][0]
+ self.queue = Queue.Queue(capacity)
+ self.free_queue = Queue.Queue()
+ self.reset_begin = threading.Event()
+ self.reset_end = threading.Event()
+
+ self.exit = False
+ self.iter_complete = False
+ self.current_batch = None
+ def prefetch_func():
+ """Thread entry"""
+ while True:
+ if self.exit:
+ break
+ if self.reset_begin.is_set():
+ self.reset_begin.clear()
+ self.data_iter.reset()
+ self.__clear_queue()
+ # notify reset over
+ self.reset_end.set()
+ try:
+ # get free data
+ try:
+ free_data = self.free_queue.get_nowait()
+ except Queue.Empty:
+ batch_data = nd.empty(self.provide_data[0][1])
+ batch_label = nd.empty(self.provide_label[0][1])
+ free_data = io.DataBatch([batch_data], [batch_label])
+ next_data = self.data_iter.next_batch(free_data)
+ self.queue.put(next_data, block=True)
+ except StopIteration:
+ self.queue.put(PrefetchDataFlag.ITER_COMPLETE, block=True)
+ # wait for new epoch
+ self.reset_begin.wait()
+
+ self.prefetch_thread = threading.Thread(target=prefetch_func, args=[])
+ self.prefetch_thread.setDaemon(True)
+ self.prefetch_thread.start()
+
+ def __del__(self):
+ self.exit = True
+ self.__clear_queue()
+ self.reset_begin.set()
+ self.prefetch_thread.join()
+
+ def __clear_queue(self):
+ """ clear the queue"""
+ while True:
+ try:
+ self.queue.get_nowait()
+ except Queue.Empty:
+ break
+
+ @property
+ def provide_data(self):
+ return self.data_iter.provide_data
+
+ @property
+ def provide_label(self):
+ return self.data_iter.provide_label
+
+ def reset(self):
+ # notify reset begining
+ self.reset_begin.set()
+ self.__clear_queue()
+ self.iter_complete = False
+ # wait for reset over
+ self.reset_end.wait()
+ self.reset_end.clear()
+
+ def iter_next(self):
+ if self.iter_complete:
+ return False
+ else:
+ # recyle first
+ if self.current_batch is not None:
+ self.free_queue.put_nowait(self.current_batch)
+ self.current_batch = None
+ d = self.queue.get(block=True)
+ if d == PrefetchDataFlag.ITER_COMPLETE:
+ self.iter_complete = True
+ return False
+ else:
+ self.current_batch = d
+ return True
+
+ def next(self):
+ if self.iter_next():
+ return self.current_batch
+ else:
+ raise StopIteration
+
+ def getdata(self):
+ return self.current_batch.data
+
+ def getlabel(self):
+ return self.current_batch.label
+
+ def getindex(self):
+ return self.current_batch.index
+
+ def getpad(self):
+ return self.current_batch.pad
+
+def PrefetchImageIter(prefetch_num = 4, **kwargs):
+ """Performs Pre-fetch for ImageIter
+
+ This iterator will create another thread to perform ``iter_next`` and then
+ store the data in queue. It potentially accelerates the data read, at the
+ cost of more memory usage.
+
+ Parameters
+ ----------
+ prefetch_num : int
+ prefetch num of io.DataBatch object
+ kwargs : ...
+ More arguments for creating ImageIter. See mx.image.ImageIter.
+ """
+ data_iter = PrefetchThreadIter(ImageIter(**kwargs), prefetch_num)
+ import atexit
+ atexit.register(lambda a : a.__del__(), data_iter)
+ return data_iter
diff --git a/python/mxnet/ndarray/utils.py b/python/mxnet/ndarray/utils.py
index 6f3b0ff9c5..d4fc3b4279 100644
--- a/python/mxnet/ndarray/utils.py
+++ b/python/mxnet/ndarray/utils.py
@@ -236,3 +236,61 @@ def save(fname, data):
mx_uint(len(handles)),
c_array(NDArrayHandle, handles),
keys))
+
+def load_from_str(data):
+ """Loads an array from string.
+
+ Parameters
+ ----------
+ data : str
+ string of ndarray
+
+ Returns
+ -------
+ NDArray
+ loaded data
+ """
+ handle = NDArrayHandle()
+ length = ctypes.c_size_t(len(data))
+ check_call(_LIB.MXNDArrayLoadFromRawBytes(c_str(data), length, ctypes.byref(handle)))
+ return NDArray(handle)
+
+def save_to_str(data):
+ """save an array to string.
+
+ Parameters
+ ----------
+ data : NDArray
+ The data to save.
+
+ Returns
+ -------
+ bytes
+ NDArray string of the data
+ Examples
+ --------
+ >>> x = mx.nd.zeros((2,3))
+ >>> y = mx.nd.save_to_str(x)
+ >>> mx.nd.load_from_str(y)
+ <NDArray 2x3 @cpu(0)>
+
+ --------
+ In [23]: a = nd.zeros((1000,1000))
+
+ In [24]: %timeit(cPickle.dumps(a))
+ 10 loops, best of 3: 22 ms per loop
+
+ In [25]: %timeit(nd.save_to_str(a))
+ 100 loops, best of 3: 2.24 ms per loop
+ """
+ handle = data.handle
+ if handle is None:
+ return ''
+ length = ctypes.c_size_t()
+ cptr = ctypes.POINTER(ctypes.c_char)()
+ check_call(_LIB.MXNDArraySaveRawBytes(handle,
+ ctypes.byref(length),
+ ctypes.byref(cptr)))
+ buf = bytes(ctypes2buffer(cptr, length.value))
+ return buf
+
diff --git a/python/mxnet/recordio.py b/python/mxnet/recordio.py
index 39f442b6ae..b37fd23cab 100644
--- a/python/mxnet/recordio.py
+++ b/python/mxnet/recordio.py
@@ -19,15 +19,17 @@
from __future__ import absolute_import
from collections import namedtuple
+import sys
import ctypes
import struct
import numbers
import numpy as np
+from . import ndarray as nd
from .base import _LIB
-from .base import RecordIOHandle
-from .base import check_call
-from .base import c_str
+from .base import RecordIOHandle, RecordIterHandle, NDArrayHandle
+from .base import check_call, ctypes2buffer, build_param_doc as _build_param_doc
+from .base import c_str, c_array, mx_uint, py_str
try:
import cv2
except ImportError:
@@ -454,3 +456,261 @@ def pack_img(header, img, quality=95, img_fmt='.jpg'):
ret, buf = cv2.imencode(img_fmt, img, encode_params)
assert ret, 'failed to encode image'
return pack(header, buf.tostring())
+# pb version of recordio
+RecordIOHead = namedtuple('HEAD', ['id', 'reserve'])
+RecordIOData = namedtuple('DATA', ['id', 'type', 'nd_value', 'str_value'])
+RecordIOExtraData = namedtuple('EXTRA_DATA', ['key', 'type', 'nd_value', 'str_value'])
+class RecordIOType(object):
+ NDARRAY = 0
+ BINARY = 1
+ STRING = 2
+
+class MXRecordIO_PB(object):
+ """ a recordio from pb-format files.
+
+ Parameters
+ ----------
+ data : list of `RecordIOData`
+ A list of input data.
+ label : list of `float`
+ A list of input label.
+ header : RecordIOHead
+ head of input record
+ extra_data : RecordIOExtraData
+ A list of input extra data.
+ """
+ def __init__(self, data, label=None, header=None, extra_data=None):
+ if data is not None:
+ assert isinstance(data, list), "Data must be list"
+ if label is not None:
+ assert isinstance(label, list), "Label must be list"
+ if extra_data is not None:
+ assert isinstance(extra_data, list), "extra_data must be list"
+ self.data = data
+ self.label = label
+ self.header = header
+ self.extra_data = extra_data
+
+class MXRecordIter(object):
+ """A python wrapper a C++ data iterator.
+
+ This iterator is the Python wrapper to native C++ RecordIter iterator.
+ When initializing `RecordIter` for example, you will get an `MXRecordIter`
+ instance to use in your Python code.
+ Calls to `next`, `reset`, etc will be delegated to the
+ underlying C++ data iterators.
+
+ Usually you don't need to interact with `MXRecordIter` directly unless you are
+ implementing your own data iterators in C++. To do that, please refer to
+ examples under the `src/io` folder.
+
+ Parameters
+ ----------
+ handle : DataIterHandle, required
+ The handle to the underlying C++ Data Iterator.
+
+ See Also
+ --------
+ src/io : The underlying C++ data iterator implementation.
+ """
+ def __init__(self, handle):
+ self.handle = handle
+
+ def __del__(self):
+ check_call(_LIB.MXRecordIterFree(self.handle))
+ self.handle = RecordIterHandle()
+
+ def __iter__(self):
+ return self
+
+ def reset(self):
+ """Reset the iterator to the begin of the data."""
+ check_call(_LIB.MXRecordIterBeforeFirst(self.handle))
+
+ def next(self):
+ if self.iter_next():
+ return MXRecordIO_PB(data = self.get_data(),
+ label = self.get_label(),
+ header = self.get_head(),
+ extra_data = self.get_extra_data())
+ else:
+ raise StopIteration
+
+ def __next__(self):
+ return self.next()
+
+ def iter_next(self):
+ """Move to the next record.
+
+ Returns
+ -------
+ boolean
+ Whether the move is successful.
+ """
+ next_res = ctypes.c_int(0)
+ check_call(_LIB.MXRecordIterNext(self.handle, ctypes.byref(next_res)))
+ return next_res.value
+
+ def get_data(self):
+ """Get data of current record.
+ """
+ num = self.get_data_num()
+ data_list = []
+ for i in xrange(num):
+ data_id = ctypes.c_uint64(0)
+ data_type = ctypes.c_int(0)
+ data_ndarray = NDArrayHandle()
+ data_string_length = mx_uint()
+ data_string = ctypes.POINTER(ctypes.c_char)()
+ check_call(_LIB.MXRecordIterGetData(self.handle,
+ ctypes.c_uint(i),
+ ctypes.byref(data_id),
+ ctypes.byref(data_type),
+ ctypes.byref(data_ndarray),
+ ctypes.byref(data_string_length),
+ ctypes.byref(data_string)))
+ nd_value = nd.NDArray(data_ndarray, False)
+ data_list.append(RecordIOData(id=data_id.value, type=data_type.value,
+ nd_value=nd_value,
+ str_value=ctypes2buffer(data_string, data_string_length.value)))
+ return data_list
+
+ def get_extra_data(self):
+ """Get extra data of current record.
+ """
+ num = self.get_extra_data_num()
+ data_list = []
+ for i in xrange(num):
+ data_key = ctypes.POINTER(ctypes.c_char)()
+ data_key_length = mx_uint()
+ data_type = ctypes.c_int(0)
+ data_ndarray = NDArrayHandle()
+ data_string = ctypes.POINTER(ctypes.c_char)()
+ data_string_length = mx_uint()
+ check_call(_LIB.MXRecordIterGetExtraData(self.handle,
+ ctypes.c_uint(i),
+ ctypes.byref(data_key),
+ ctypes.byref(data_key_length),
+ ctypes.byref(data_type),
+ ctypes.byref(data_ndarray),
+ ctypes.byref(data_string),
+ ctypes.byref(data_string_length)))
+ nd_value = nd.NDArray(data_ndarray, False)
+ data_list.append(RecordIOExtraData(key=ctypes2buffer(data_key, data_key_length.value),
+ type=data_type.value,
+ nd_value=nd_value,
+ str_value=ctypes2buffer(data_string, data_string_length.value)))
+ return data_list
+
+ def get_data_num(self):
+ """Get data num of current record.
+ """
+ num = ctypes.c_uint(0)
+ check_call(_LIB.MXRecordIterGetDataNum(self.handle, ctypes.byref(num)))
+ return num.value
+
+ def get_extra_data_num(self):
+ """Get extra data num of current record.
+ """
+ num = ctypes.c_uint(0)
+ check_call(_LIB.MXRecordIterGetExtraDataNum(self.handle, ctypes.byref(num)))
+ return num.value
+
+ def get_label(self):
+ """Get label of the current record.
+ """
+ label_size = ctypes.c_uint(0)
+ label_data = ctypes.POINTER(ctypes.c_float)()
+ check_call(_LIB.MXRecordIterGetLabel(self.handle,
+ ctypes.byref(label_data),
+ ctypes.byref(label_size)))
+ label_list = [label_data[i] for i in xrange(label_size.value)]
+ return label_list
+
+ def get_head(self):
+ """Get head of the current record.
+ """
+ record_id = ctypes.c_uint64(0)
+ reserve = ctypes.c_uint64(0)
+ check_call(_LIB.MXRecordIterGetHead(self.handle,
+ ctypes.byref(record_id), ctypes.byref(reserve)))
+ return RecordIOHead(id=record_id.value, reserve=reserve.value)
+
+def _make_record_iterator():
+ """Create an io iterator by handle."""
+ name = ctypes.c_char_p()
+ desc = ctypes.c_char_p()
+ num_args = mx_uint()
+ arg_names = ctypes.POINTER(ctypes.c_char_p)()
+ arg_types = ctypes.POINTER(ctypes.c_char_p)()
+ arg_descs = ctypes.POINTER(ctypes.c_char_p)()
+
+ check_call(_LIB.MXRecordIterGetIterInfo( \
+ ctypes.byref(name), ctypes.byref(desc), \
+ ctypes.byref(num_args), \
+ ctypes.byref(arg_names), \
+ ctypes.byref(arg_types), \
+ ctypes.byref(arg_descs)))
+ if name.value is None:
+ return None
+ iter_name = py_str(name.value)
+
+ narg = int(num_args.value)
+ param_str = _build_param_doc(
+ [py_str(arg_names[i]) for i in range(narg)],
+ [py_str(arg_types[i]) for i in range(narg)],
+ [py_str(arg_descs[i]) for i in range(narg)])
+
+ doc_str = ('%s\n\n' +
+ '%s\n' +
+ 'Returns\n' +
+ '-------\n' +
+ 'MXRecordIter\n'+
+ ' The result iterator.')
+ doc_str = doc_str % (desc.value, param_str)
+
+ def creator(*args, **kwargs):
+ """Create an iterator.
+ The parameters listed below can be passed in as keyword arguments.
+
+ Parameters
+ ----------
+ kwargs : ...
+ arguments for creating MXRecordIter. See MXRecordIterCreate.
+
+ Returns
+ -------
+ MXRecordIter
+ The resulting data iterator.
+ """
+ param_keys = []
+ param_vals = []
+
+ for k, val in kwargs.items():
+ param_keys.append(c_str(k))
+ param_vals.append(c_str(str(val)))
+ # create atomic symbol
+ param_keys = c_array(ctypes.c_char_p, param_keys)
+ param_vals = c_array(ctypes.c_char_p, param_vals)
+ iter_handle = RecordIterHandle()
+ check_call(_LIB.MXRecordIterCreate(
+ mx_uint(len(param_keys)),
+ param_keys, param_vals,
+ ctypes.byref(iter_handle)))
+
+ if len(args):
+ raise TypeError('%s can only accept keyword arguments' % iter_name)
+
+ return MXRecordIter(iter_handle)
+
+ creator.__name__ = iter_name
+ creator.__doc__ = doc_str
+ return creator
+
+def _init_record_module():
+ module_obj = sys.modules[__name__]
+ record_iter = _make_record_iterator()
+ if record_iter is not None:
+ setattr(module_obj, record_iter.__name__, record_iter)
+
+_init_record_module()
diff --git a/src/c_api/c_api.cc b/src/c_api/c_api.cc
index 1d348a5b40..95d8e57622 100644
--- a/src/c_api/c_api.cc
+++ b/src/c_api/c_api.cc
@@ -44,6 +44,7 @@
#include "./c_api_common.h"
#include "../operator/custom/custom-inl.h"
#include "../engine/profiler.h"
+#include "../io/iter_recordio.h"
using namespace mxnet;
@@ -1228,3 +1229,202 @@ int MXRtcCudaKernelCall(CudaKernelHandle handle, int dev_id, void** args,
#endif
API_END();
}
+
+// RecordIter
+int MXRecordIterGetIterInfo(const char **name,
+ const char **description,
+ mx_uint *num_args,
+ const char ***arg_names,
+ const char ***arg_type_infos,
+ const char ***arg_descriptions) {
+#if PB_FORMAT_REC
+ auto &vec = dmlc::Registry<io::RecordIteratorReg>::List();
+ CHECK(vec.size() == 1) << "RecordIter register error";
+ auto e = vec[0];
+ return MXAPIGetFunctionRegInfo(e, name, description, num_args,
+ arg_names, arg_type_infos, arg_descriptions,
+ NULL);
+#else
+ API_BEGIN();
+ API_END();
+#endif
+}
+
+int MXRecordIterCreate(mx_uint num_param,
+ const char **keys,
+ const char **vals,
+ RecordIterHandle *out) {
+#if PB_FORMAT_REC
+ IIterator<io::RecordInst> *iter = nullptr;
+ API_BEGIN();
+ auto &vec = dmlc::Registry<io::RecordIteratorReg>::List();
+ CHECK(vec.size() == 1) << "RecordIter register error";
+ auto e = vec[0];
+ iter = e->body();
+ std::vector<std::pair<std::string, std::string> > kwargs;
+ for (mx_uint i = 0; i < num_param; ++i) {
+ kwargs.push_back({std::string(keys[i]), std::string(vals[i])});
+ }
+ iter->Init(kwargs);
+ *out = iter;
+ API_END_HANDLE_ERROR(delete iter);
+#else
+ API_BEGIN();
+ LOG(FATAL) << "Need to compile with USE_PB_FORMAT_REC=1 for MXRecordIter.";
+ API_END();
+#endif
+}
+
+int MXRecordIterFree(RecordIterHandle handle) {
+ API_BEGIN();
+#if PB_FORMAT_REC
+ delete static_cast<IIterator<io::RecordInst> *>(handle);
+#else
+ LOG(FATAL) << "Need to compile with USE_PB_FORMAT_REC=1 for MXRecordIter.";
+#endif
+ API_END();
+}
+
+int MXRecordIterBeforeFirst(RecordIterHandle handle) {
+ API_BEGIN();
+#if PB_FORMAT_REC
+ static_cast<IIterator<io::RecordInst>* >(handle)->BeforeFirst();
+#else
+ LOG(FATAL) << "Need to compile with USE_PB_FORMAT_REC=1 for MXRecordIter.";
+#endif
+ API_END();
+}
+
+int MXRecordIterNext(RecordIterHandle handle, int *out) {
+ API_BEGIN();
+#if PB_FORMAT_REC
+ *out = static_cast<IIterator<io::RecordInst>* >(handle)->Next();
+#else
+ LOG(FATAL) << "Need to compile with USE_PB_FORMAT_REC=1 for MXRecordIter.";
+#endif
+ API_END();
+}
+
+int MXRecordIterGetDataNum(RecordIterHandle handle,
+ mx_uint *out) {
+ API_BEGIN();
+#if PB_FORMAT_REC
+ const io::RecordInst& record = static_cast<IIterator<io::RecordInst>* >(handle)->Value();
+ *out = static_cast<mx_uint>(record.data.size());
+#else
+ LOG(FATAL) << "Need to compile with USE_PB_FORMAT_REC=1 for MXRecordIter.";
+#endif
+ API_END();
+}
+
+int MXRecordIterGetData(RecordIterHandle handle,
+ mx_uint index,
+ uint64_t *out_id,
+ int *out_type,
+ NDArrayHandle *out_ndarray,
+ mx_uint *out_string_len,
+ const char **out_string) {
+ NDArray *ptr = nullptr;
+ API_BEGIN();
+#if PB_FORMAT_REC
+ const auto& record = static_cast<IIterator<io::RecordInst>* >(handle)->Value().data;
+ *out_id = record[index].id;
+ *out_type = record[index].type;
+ const char * tmp_str = nullptr;
+ mx_uint tmp_str_len = 0;
+ switch (record[index].type) {
+ case io::record_data_type::kNdarray:
+ case io::record_data_type::kBinary:
+ ptr = new NDArray();
+ *ptr = record[index].value.nd;
+ break;
+ case io::record_data_type::kString:
+ tmp_str = record[index].value.str.c_str();
+ tmp_str_len = static_cast<mx_uint>(record[index].value.str.length());
+ default:
+ break;
+ }
+ *out_ndarray = ptr;
+ *out_string = tmp_str;
+ *out_string_len = tmp_str_len;
+#else
+ LOG(FATAL) << "Need to compile with USE_PB_FORMAT_REC=1 for MXRecordIter.";
+#endif
+ API_END_HANDLE_ERROR(delete ptr);
+}
+
+int MXRecordIterGetExtraDataNum(RecordIterHandle handle,
+ mx_uint *out) {
+ API_BEGIN();
+#if PB_FORMAT_REC
+ const io::RecordInst& record = static_cast<IIterator<io::RecordInst>* >(handle)->Value();
+ *out = static_cast<mx_uint>(record.extra.size());
+#else
+ LOG(FATAL) << "Need to compile with USE_PB_FORMAT_REC=1 for MXRecordIter.";
+#endif
+ API_END();
+}
+
+int MXRecordIterGetExtraData(RecordIterHandle handle,
+ mx_uint index,
+ const char **out_key,
+ mx_uint *out_key_len,
+ int *out_type,
+ NDArrayHandle *out_ndarray,
+ const char **out_string,
+ mx_uint *out_string_len) {
+ NDArray *ptr = nullptr;
+ API_BEGIN();
+#if PB_FORMAT_REC
+ const auto& record = static_cast<IIterator<io::RecordInst>* >(handle)->Value().extra;
+ *out_key = record[index].key.c_str();
+ *out_key_len = static_cast<mx_uint>(record[index].key.length());
+ *out_type = record[index].type;
+ const char * tmp_str = nullptr;
+ mx_uint tmp_str_len = 0;
+ switch (record[index].type) {
+ case io::record_data_type::kNdarray:
+ case io::record_data_type::kBinary:
+ ptr = new NDArray();
+ *ptr = record[index].value.nd;
+ break;
+ case io::record_data_type::kString:
+ tmp_str = record[index].value.str.c_str();
+ tmp_str_len = static_cast<mx_uint>(record[index].value.str.length());
+ break;
+ default:
+ break;
+ }
+ *out_ndarray = ptr;
+ *out_string = tmp_str;
+ *out_string_len = tmp_str_len;
+#else
+ LOG(FATAL) << "Need to compile with USE_PB_FORMAT_REC=1 for MXRecordIter.";
+#endif
+ API_END_HANDLE_ERROR(delete ptr);
+}
+
+int MXRecordIterGetLabel(RecordIterHandle handle, mx_float **out, mx_uint *out_size) {
+ API_BEGIN();
+#if PB_FORMAT_REC
+ const io::RecordInst& record = static_cast<IIterator<io::RecordInst>* >(handle)->Value();
+ *out_size = static_cast<mx_uint>(record.label.size());
+ *out = const_cast<float*>(dmlc::BeginPtr(record.label));
+#else
+ LOG(FATAL) << "Need to compile with USE_PB_FORMAT_REC=1 for MXRecordIter.";
+#endif
+ API_END();
+}
+
+int MXRecordIterGetHead(RecordIterHandle handle, uint64_t *out_id, uint64_t *out_reserve) {
+ API_BEGIN();
+#if PB_FORMAT_REC
+ const io::RecordInst& record = static_cast<IIterator<io::RecordInst>* >(handle)->Value();
+ *out_id = record.head.id();
+ *out_reserve = record.head.reserve();
+#else
+ LOG(FATAL) << "Need to compile with USE_PB_FORMAT_REC=1 for MXRecordIter.";
+#endif
+ API_END();
+}
+
diff --git a/src/common/cuda_utils.h b/src/common/cuda_utils.h
index c135ff8a1d..2e5fe0dda5 100644
--- a/src/common/cuda_utils.h
+++ b/src/common/cuda_utils.h
@@ -87,6 +87,7 @@ inline const char* CublasGetErrorString(cublasStatus_t error) {
return "Unknown cuBLAS status";
}
+#if MXNET_USE_CUSOLVER == 1
/*!
* \brief Get string representation of cuSOLVER errors.
* \param error The error.
@@ -115,6 +116,7 @@ inline const char* CusolverGetErrorString(cusolverStatus_t error) {
}
return "Unknown cuSOLVER status";
}
+#endif
/*!
* \brief Get string representation of cuRAND errors.
diff --git a/src/io/iter_recordio.cc b/src/io/iter_recordio.cc
new file mode 100644
index 0000000000..7eec1134cd
--- /dev/null
+++ b/src/io/iter_recordio.cc
@@ -0,0 +1,37 @@
+/*!
+ * Copyright (c) 2017 by shuqian.qu@hobot.cc
+ * \file iter_recordio.cc
+ * \brief record iter
+ */
+#include "./iter_recordio.h"
+
+#if PB_FORMAT_REC
+// Registers
+namespace dmlc {
+DMLC_REGISTRY_ENABLE(::mxnet::io::RecordIteratorReg);
+} // namespace dmlc
+
+namespace mxnet {
+namespace io {
+DMLC_REGISTER_PARAMETER(RecParserParam);
+DMLC_REGISTER_PARAMETER(RecordDataInstParam);
+
+MXNET_REGISTER_RECORD_ITER(RecordIter)
+.describe(R"code(Iterating on RecordIO pb-format files
+
+Read one record from RecordIO files with a rich of data load options.
+
+One can use ``tools/im2rec.py --pb-format=1`` to pack individual files
+into RecordIO files.
+
+)code" ADD_FILELINE)
+.add_arguments(RecParserParam::__FIELDS__())
+.add_arguments(RecordDataInstParam::__FIELDS__())
+.set_body([]() {
+ return new RecordIter();
+ });
+
+} // namespace io
+} // namespace mxnet
+
+#endif // PB_FORMAT_REC
diff --git a/src/io/iter_recordio.h b/src/io/iter_recordio.h
new file mode 100644
index 0000000000..7ffe54c0aa
--- /dev/null
+++ b/src/io/iter_recordio.h
@@ -0,0 +1,470 @@
+/*!
+ * Copyright (c) 2017 by shuqian.qu@hobot.cc
+ * \file iter_recordio.h
+ * \brief Interface of recordio iter
+ */
+
+#ifndef MXNET_IO_ITER_RECORDIO_H_
+#define MXNET_IO_ITER_RECORDIO_H_
+
+#include <mxnet/io.h>
+#include <mxnet/base.h>
+#include <dmlc/recordio.h>
+#include <dmlc/memory_io.h>
+#include <dmlc/threadediter.h>
+#include <dmlc/input_split_shuffle.h>
+#include <utility>
+#include <vector>
+#include <string>
+#include <algorithm>
+#include "../common/utils.h"
+#include "./iter_threaded_input_split.h"
+
+#if PB_FORMAT_REC
+#include "./recordio.pb.h"
+namespace mxnet {
+namespace io {
+
+namespace record_data_type {
+enum RecordDataType {kNdarray, kBinary, kString};
+}
+
+struct RecordInst {
+ // value
+ struct TValue {
+ NDArray nd;
+ std::string str;
+ };
+ // record data
+ struct RData {
+ uint64_t id;
+ int type;
+ TValue value;
+ };
+ // extra data
+ struct ExData {
+ std::string key;
+ int type;
+ TValue value;
+ };
+
+ RecordHead head;
+ std::vector<RData> data;
+ std::vector<float> label;
+ std::vector<ExData> extra;
+};
+
+/*!
+ * \brief a list
+ */
+class RecordInstVector {
+ public:
+ /*! \brief return the number of record */
+ inline size_t Size(void) const {
+ return record_list_.size();
+ }
+ // instance
+ /* \brief get the i-th record */
+ inline RecordInst operator[](size_t i) const {
+ return record_list_[i];
+ }
+ /* \brief get the last (label, example) pair */
+ inline RecordInst Back() const {
+ return (*this)[Size() - 1];
+ }
+
+ inline void Clear(void) {
+ record_list_.clear();
+ }
+
+ inline void Push(const RecordInst& rec) {
+ record_list_.push_back(rec);
+ }
+
+ std::vector<RecordInst> record_list_;
+};
+
+// Define record parser parameters
+struct RecParserParam : public dmlc::Parameter<RecParserParam> {
+ /*! \brief the ratio of data sampling*/
+ float sampling_ratio;
+ DMLC_DECLARE_PARAMETER(RecParserParam) {
+ DMLC_DECLARE_FIELD(sampling_ratio).set_default(1.0f)
+ .describe("the ratio of data sampling");
+ }
+};
+
+// parser to parse recordio
+class RecordIOParser {
+ public:
+ // initialize the parser
+ inline void Init(const std::vector<std::pair<std::string, std::string> >& kwargs,
+ int parse_thread);
+
+ // set record to the head
+ inline void BeforeFirst(void) { }
+ // parse next set of records, return an array of
+ // instance vector to the user
+ inline bool ParseNext(RecordInstVector *out_vec,
+ dmlc::InputSplit::Blob chunk,
+ int tid);
+
+ private:
+ // magic number to see prng
+ static const int kRandMagic = 111;
+ /*! \brief parameters */
+ RecParserParam param_;
+ /*! \brief random samplers */
+ std::vector<std::unique_ptr<common::RANDOM_ENGINE> > prnds_;
+};
+
+inline void LoadRecTypeValue(RecordInst::TValue* value,
+ const std::string& origin, int type) {
+ switch (type) {
+ case record_data_type::kNdarray:
+ {
+ // NDArray value
+ auto val = origin;
+ dmlc::MemoryStringStream strm(&val); // NOLINT(*)
+ if (!value->nd.Load(&strm)) {
+ LOG(FATAL) << "Invalid NDArray serialization format!";
+ }
+ }
+ break;
+ case record_data_type::kBinary:
+ {
+ // string to NDArray type
+ value->nd = NDArray(mshadow::Shape1(origin.size()),
+ Context::CPU(),
+ false,
+ mshadow::kUint8);
+ auto nd_blob = value->nd.data();
+ memcpy(nd_blob.dptr_, origin.c_str(), origin.size());
+ }
+ break;
+ case record_data_type::kString:
+ {
+ value->str = origin;
+ }
+ break;
+ default:
+ {
+ LOG(FATAL) << "read data type error:" << type;
+ }
+ break;
+ }
+ return;
+}
+
+inline void RecordIOParser::Init(
+ const std::vector<std::pair<std::string, std::string> >& kwargs,
+ int parse_thread) {
+ // initialize parameter
+ param_.InitAllowUnknown(kwargs);
+ for (int i = 0; i < parse_thread; ++i) {
+ prnds_.emplace_back(new common::RANDOM_ENGINE((i + 1) * kRandMagic));
+ }
+}
+
+inline bool RecordIOParser::ParseNext(RecordInstVector *out_data,
+ dmlc::InputSplit::Blob chunk,
+ int tid) {
+ dmlc::RecordIOChunkReader reader(chunk);
+ auto& out = *out_data;
+ out.Clear();
+ dmlc::InputSplit::Blob blob;
+ while (reader.NextRecord(&blob)) {
+ if (param_.sampling_ratio < 1.0f) {
+ std::uniform_real_distribution<float> dis(0, 1);
+ // drop all samples in cat
+ if (dis(*prnds_[tid]) > param_.sampling_ratio) continue;
+ }
+ if (blob.size <= 0) continue;
+ RecordInst rec;
+ RecordUnit pb_record;
+ pb_record.ParseFromString(std::string(static_cast<const char*>(blob.dptr), blob.size));
+ if (!pb_record.has_head() || !pb_record.has_body()) continue;
+ rec.head = pb_record.head();
+ const auto& body = pb_record.body();
+ // data
+ rec.data.resize(body.data_size());
+ for (int i = 0; i < body.data_size(); i++) {
+ const auto& data_i = body.data(i);
+ rec.data[i].id = data_i.id();
+ rec.data[i].type = data_i.type();
+ LoadRecTypeValue(&rec.data[i].value, data_i.value(), data_i.type());
+ }
+ // label
+ for (int i = 0; i < body.label_size(); i++) {
+ rec.label.push_back(body.label(i));
+ }
+ // extra
+ rec.extra.resize(body.extra_size());
+ for (int i = 0; i < body.extra_size(); i++) {
+ const auto& extra_i = body.extra(i);
+ rec.extra[i].key = extra_i.key();
+ rec.extra[i].type = extra_i.type();
+ LoadRecTypeValue(&rec.extra[i].value, extra_i.value(), extra_i.type());
+ }
+ out.Push(rec);
+ }
+ return true;
+}
+
+// RecordDataInstIter parameters
+struct RecordDataInstParam : public dmlc::Parameter<RecordDataInstParam> {
+ /*! \brief whether to do shuffle */
+ bool shuffle;
+ /*! \brief random seed */
+ int seed;
+ /*! \brief whether to remain silent */
+ bool verbose;
+ /*! \brief path to recordio */
+ std::string path_imgrec;
+ /*! \brief number of threads */
+ int preprocess_threads;
+ /*! \brief whether to read by chunk*/
+ bool chunk_read;
+ /*! \brief size of data chunk */
+ int chunk_size;
+ /*! \brief partition the data into multiple parts */
+ int num_parts;
+ /*! \brief the index of the part will read*/
+ int part_index;
+ /*! \brief the size of a shuffle chunk*/
+ size_t shuffle_chunk_size;
+ /*! \brief the seed for chunk shuffling*/
+ int shuffle_chunk_seed;
+ /*! \brief source chunk's queue capacity*/
+ int source_chunk_queue_capacity;
+ /*! \brief record chunk's queue capacity*/
+ int record_chunk_queue_capacity;
+ // declare parameters
+ DMLC_DECLARE_PARAMETER(RecordDataInstParam) {
+ DMLC_DECLARE_FIELD(shuffle).set_default(false)
+ .describe("Augmentation Param: Whether to shuffle data.");
+ DMLC_DECLARE_FIELD(seed).set_default(0)
+ .describe("Augmentation Param: Random Seed.");
+ DMLC_DECLARE_FIELD(verbose).set_default(true)
+ .describe("Auxiliary Param: Whether to output information.");
+ DMLC_DECLARE_FIELD(path_imgrec).set_default("")
+ .describe("Dataset Param: Path to record file.");
+ DMLC_DECLARE_FIELD(preprocess_threads).set_lower_bound(1).set_default(2)
+ .describe("Backend Param: Number of thread to do preprocessing.");
+ DMLC_DECLARE_FIELD(chunk_size).set_lower_bound(4).set_default(16)
+ .describe("Backend Param: chunk_size.");
+ DMLC_DECLARE_FIELD(chunk_read).set_default(true)
+ .describe("Backend Param: whether to read by chunk.");
+ DMLC_DECLARE_FIELD(num_parts).set_default(1)
+ .describe("partition the data into multiple parts");
+ DMLC_DECLARE_FIELD(part_index).set_default(0)
+ .describe("the index of the part will read");
+ DMLC_DECLARE_FIELD(shuffle_chunk_size).set_default(0)
+ .describe("the size(MB) of the shuffle block, "
+ "it is recommended to use 16 or more when using hdfs");
+ DMLC_DECLARE_FIELD(shuffle_chunk_seed).set_default(0)
+ .describe("the seed for block shuffling");
+ DMLC_DECLARE_FIELD(source_chunk_queue_capacity).set_default(2)
+ .describe("source chunk's queue capacity");
+ DMLC_DECLARE_FIELD(record_chunk_queue_capacity).set_default(2)
+ .describe("record chunk's queue capacity");
+ }
+};
+
+template <typename ParserType, typename InstVecType, typename InstType = DataInst>
+class RecordDataInstIter : public IIterator<InstType> {
+ public:
+ RecordDataInstIter(): inst_ptr_(0), data_(nullptr), iter_(nullptr) { }
+ // destructor
+ virtual ~RecordDataInstIter(void) {
+ delete iter_;
+ iter_ = nullptr;
+ delete data_;
+ data_ = nullptr;
+ }
+ // constructor
+ virtual void Init(const std::vector<std::pair<std::string, std::string> >& kwargs);
+ // before first
+ virtual void BeforeFirst(void) {
+ iter_->BeforeFirst();
+ inst_order_.clear();
+ inst_ptr_ = 0;
+ }
+
+ virtual bool Next(void);
+
+ virtual const InstType &Value(void) const {
+ return out_;
+ }
+
+ private:
+ // random magic
+ static const int kRandMagic = 11111;
+ // output instance
+ InstType out_;
+ // data ptr
+ size_t inst_ptr_;
+ // internal instance order
+ std::vector<unsigned> inst_order_;
+ // data
+ InstVecType* data_;
+ // internal parser
+ ParserType parser_;
+ /*! \brief data source */
+ std::unique_ptr<dmlc::InputSplit> source_;
+ // backend thread
+ dmlc::MultiThreadedIter<InstVecType, CommonChunkType>* iter_;
+ // parameters
+ RecordDataInstParam param_;
+ // random number generator
+ common::RANDOM_ENGINE rnd_;
+};
+
+template <typename ParserType, typename InstVecType, typename InstType>
+void RecordDataInstIter<ParserType, InstVecType, InstType>::Init(
+ const std::vector<std::pair<std::string, std::string> >& kwargs) {
+ const auto& unknown_params = param_.InitAllowUnknown(kwargs);
+ for (size_t index = 0; index < unknown_params.size(); index++) {
+ CHECK(unknown_params[index].first != "shuffle_seed") <<
+ "shuffle_seed param is deprecated, please use shuffle_chunk_seed instead!";
+ }
+ CHECK(param_.path_imgrec.length() != 0)
+ << "RecordDataInstIter: must specify path_imgrec";
+ // preprocess thread
+#if defined(_OPENMP)
+ int maxthread = 0;
+ int threadget = 0;
+ #pragma omp parallel
+ {
+ // be conservative, set number of real cores
+ maxthread = std::max(omp_get_num_procs() - 1, 1);
+ }
+ param_.preprocess_threads = std::min(maxthread, param_.preprocess_threads);
+ #pragma omp parallel num_threads(param_.preprocess_threads)
+ {
+ threadget = omp_get_num_threads();
+ }
+ param_.preprocess_threads = threadget;
+#endif
+
+ // use the kwarg to init parser
+ parser_.Init(kwargs, param_.preprocess_threads);
+ // input split
+ source_.reset(dmlc::InputSplit::Create(
+ param_.path_imgrec.c_str(), param_.part_index,
+ param_.num_parts, "recordio"));
+ // use 64 MB chunk when possible
+ if (param_.shuffle_chunk_size > 0 && param_.shuffle) {
+ if (param_.shuffle_chunk_size > 4096) {
+ LOG(INFO) << "Chunk size: " << param_.shuffle_chunk_size
+ << " MB which is larger than 4096 MB, please set "
+ "smaller chunk size";
+ }
+ if (param_.shuffle_chunk_size < 4) {
+ LOG(INFO) << "Chunk size: " << param_.shuffle_chunk_size
+ << " MB which is less than 4 MB, please set "
+ "larger chunk size";
+ }
+
+ // 1.1 ratio is for a bit more shuffle parts to avoid boundary issue
+ unsigned num_shuffle_parts =
+ std::ceil(source_->GetTotalSize() * 1.1 /
+ (param_.num_parts * (param_.shuffle_chunk_size << 20UL)));
+ if (num_shuffle_parts > 1) {
+ source_.reset(dmlc::InputSplitShuffle::Create(
+ param_.path_imgrec.c_str(), param_.part_index,
+ param_.num_parts, "recordio", num_shuffle_parts, param_.shuffle_chunk_seed));
+ }
+ source_->HintChunkSize(param_.shuffle_chunk_size << 20UL);
+ } else {
+ // use 4 MB chunk when possible
+ source_->HintChunkSize(param_.chunk_size << 20UL);
+ }
+
+ iter_ = new dmlc::MultiThreadedIter<InstVecType, CommonChunkType>(
+ CreateInputSplitThreadedIter(source_.get(),
+ param_.source_chunk_queue_capacity,
+ param_.chunk_read),
+ param_.preprocess_threads,
+ param_.record_chunk_queue_capacity);
+ // init thread iter
+ iter_->Init([this](InstVecType **dptr, CommonChunkType* chunk, int tid) {
+ if (*dptr == nullptr) {
+ *dptr = new InstVecType();
+ }
+ dmlc::InputSplit::Blob chunk_blob;
+ chunk_blob.dptr = &(*chunk)[0];
+ chunk_blob.size = chunk->size();
+ return parser_.ParseNext(*dptr, chunk_blob, tid);
+ },
+ [this]() { parser_.BeforeFirst(); });
+ inst_ptr_ = 0;
+ rnd_.seed(kRandMagic + param_.seed);
+ if (param_.verbose) {
+ LOG(INFO) << "RecordDataInstIter: " << param_.path_imgrec
+ << ", use " << param_.preprocess_threads << " threads for decoding.. ";
+ }
+}
+
+template <typename ParserType, typename InstVecType, typename InstType>
+bool RecordDataInstIter<ParserType, InstVecType, InstType>::Next(void) {
+ while (true) {
+ if (inst_ptr_ < inst_order_.size()) {
+ unsigned p = inst_order_[inst_ptr_];
+ out_ = (*data_)[p];
+ ++inst_ptr_;
+ return true;
+ } else {
+ if (data_ != nullptr) iter_->Recycle(&data_);
+ if (!iter_->Next(&data_)) return false;
+ inst_order_.clear();
+ for (unsigned i = 0; i < data_->Size(); ++i) {
+ inst_order_.push_back(i);
+ }
+ // shuffle instance order if needed
+ if (param_.shuffle != 0) {
+ std::shuffle(inst_order_.begin(), inst_order_.end(), rnd_);
+ }
+ inst_ptr_ = 0;
+ }
+ }
+ return false;
+}
+
+// iterator on recordio
+class RecordIter : public RecordDataInstIter<RecordIOParser,
+ RecordInstVector, RecordInst> {};
+
+/*! \brief typedef the factory function of record iterator */
+typedef std::function<IIterator<RecordInst> *()> RecordIteratorFactory;
+/*!
+ * \brief Registry entry for record iterator factory functions.
+ */
+struct RecordIteratorReg
+ : public dmlc::FunctionRegEntryBase<RecordIteratorReg,
+ RecordIteratorFactory> {
+};
+//--------------------------------------------------------------
+// The following part are API Registration of Iterators
+//--------------------------------------------------------------
+/*!
+ * \brief Macro to register Iterators
+ *
+ * \code
+ * // example of registering a record iterator
+ * MXNET_REGISTER_RECORD_ITER(RecordIter)
+ * .describe("record data iterator")
+ * .set_body([]() {
+ * return new RecordIter();
+ * });
+ * \endcode
+ */
+#define MXNET_REGISTER_RECORD_ITER(name) \
+ DMLC_REGISTRY_REGISTER(::mxnet::io::RecordIteratorReg, RecordIteratorReg, name)
+
+} // namespace io
+} // namespace mxnet
+
+#endif // PB_FORMAT_REC
+
+#endif // MXNET_IO_ITER_RECORDIO_H_
diff --git a/src/io/iter_threaded_input_split.cc b/src/io/iter_threaded_input_split.cc
new file mode 100644
index 0000000000..4769595be3
--- /dev/null
+++ b/src/io/iter_threaded_input_split.cc
@@ -0,0 +1,33 @@
+// Copyright by Contributors
+
+#include <dmlc/io.h>
+#include <dmlc/threadediter.h>
+#include "./iter_threaded_input_split.h"
+
+namespace mxnet {
+namespace io {
+dmlc::ThreadedIter<CommonChunkType>* CreateInputSplitThreadedIter(
+ dmlc::InputSplit* base,
+ size_t max_capacity,
+ bool ChunkMode) {
+ CHECK(base != nullptr);
+ dmlc::ThreadedIter<CommonChunkType>* chunk_iter =
+ new dmlc::ThreadedIter<CommonChunkType>(max_capacity);
+ chunk_iter->Init([base, ChunkMode](CommonChunkType **dptr) {
+ dmlc::InputSplit::Blob blob;
+ bool base_end = ChunkMode ? (!base->NextChunk(&blob)) :
+ !base->NextRecord(&blob);
+ if (base_end) return false;
+ if (*dptr == nullptr) {
+ *dptr = new CommonChunkType();
+ }
+ (*dptr)->assign(static_cast<char*>(blob.dptr),
+ static_cast<char*>(blob.dptr) + blob.size);
+ return true;
+ },
+ [base](void) { base->BeforeFirst(); });
+ return chunk_iter;
+}
+
+} // namespace io
+} // namespace mxnet
diff --git a/src/io/iter_threaded_input_split.h b/src/io/iter_threaded_input_split.h
new file mode 100644
index 0000000000..cba8844df5
--- /dev/null
+++ b/src/io/iter_threaded_input_split.h
@@ -0,0 +1,36 @@
+/*!
+ * Copyright (c) 2015 by Contributors
+ * \file iter_threaded_input_split.h
+ * \brief
+ * \author Shuqian Qu
+ */
+#ifndef MXNET_IO_ITER_THREADED_INPUT_SPLIT_H_
+#define MXNET_IO_ITER_THREADED_INPUT_SPLIT_H_
+#include <vector>
+namespace dmlc {
+template <typename DType>
+class ThreadedIter;
+class InputSplit;
+}
+
+namespace mxnet {
+namespace io {
+/*! \brief common chunk type */
+typedef std::vector<char> CommonChunkType;
+
+/*! \brief create input split Threadediter
+ * with ThreadedIter's Next and Recycle iterface, you can consume
+ * InputSplit's data in multi thread
+ *
+ * NOTE: different from ThreadedInputSplit, which can only use in single
+ * thread
+ *
+ * \param CommonChunkType common chunk type which contains data buffer
+ */
+dmlc::ThreadedIter<CommonChunkType>* CreateInputSplitThreadedIter(
+ dmlc::InputSplit* base,
+ size_t max_capacity = 16, bool ChunkMode = true);
+
+} // namespace io
+} // namespace mxnet
+#endif // MXNET_IO_ITER_THREADED_INPUT_SPLIT_H_
diff --git a/src/io/recordio.proto b/src/io/recordio.proto
new file mode 100644
index 0000000000..755c1ce09d
--- /dev/null
+++ b/src/io/recordio.proto
@@ -0,0 +1,45 @@
+/**
+ * Copyright (c) 2017 by shuqian.qu@hobot.cc
+ */
+package mxnet.io;
+
+message RecordUnit {
+ // head
+ optional RecordHead head = 1;
+ // body
+ optional RecordBody body = 2;
+}
+
+message RecordHead {
+ // unique index
+ optional uint64 id = 1;
+ // reserve
+ optional uint64 reserve = 2;
+ // version
+ optional float version = 3;
+}
+
+message RecordBody {
+ repeated RecordData data = 1;
+ repeated float label = 2;
+ repeated ExtraData extra = 3;
+}
+
+/*!
+ type: 0 -> NDArray.
+ 1 -> binary data, return NDArray which shape is (len(value),)
+ type is uint8
+ 2 -> text data, return string in python.
+*/
+message RecordData {
+ // unique index
+ optional uint64 id = 1;
+ optional int32 type = 2;
+ optional bytes value = 3;
+}
+
+message ExtraData {
+ optional bytes key = 1;
+ optional int32 type = 2;
+ optional bytes value = 3;
+}
diff --git a/src/operator/linalg_impl.h b/src/operator/linalg_impl.h
index 0987131068..4081f6af9b 100644
--- a/src/operator/linalg_impl.h
+++ b/src/operator/linalg_impl.h
@@ -869,7 +869,7 @@ int linalg_gelqf_workspace_query<cpu, DType>(const Tensor<cpu, 2, DType>& A, \
LINALG_CPU_GELQF_WORKSPACE_QUERY(s, float)
LINALG_CPU_GELQF_WORKSPACE_QUERY(d, double)
-#ifdef __CUDACC__
+#if defined(__CUDACC__) && MXNET_USE_CUSOLVER == 1
#define LINALG_GPU_GELQF(fname, DType) \
template<> inline \
diff --git a/tools/im2rec.py b/tools/im2rec.py
index ec6de19694..73764803e2 100644
--- a/tools/im2rec.py
+++ b/tools/im2rec.py
@@ -116,15 +116,25 @@ def image_encode(args, i, item, q_out):
fullpath = os.path.join(args.root, item[1])
if len(item) > 3 and args.pack_label:
- header = mx.recordio.IRHeader(0, item[2:], item[0], 0)
+ labels = item[2:]
else:
- header = mx.recordio.IRHeader(0, item[2], item[0], 0)
+ labels = item[2]
+ if not args.pb_format:
+ header = mx.recordio.IRHeader(0, labels, item[0], 0)
+ else:
+ import recordio_pb_pack as rec_pb
+ header = rec_pb.pack_head(id=item[0], reserve=0, version=1.0)
if args.pass_through:
try:
with open(fullpath, 'rb') as fin:
img = fin.read()
- s = mx.recordio.pack(header, img)
+ if not args.pb_format:
+ s = mx.recordio.pack(header, img)
+ else:
+ import recordio_pb_pack as rec_pb
+ data = rec_pb.pack_string_data(id=item[0], data=img)
+ s = rec_pb.pack(header, data, labels)
q_out.put((i, s, item))
except Exception as e:
traceback.print_exc()
@@ -158,7 +168,14 @@ def image_encode(args, i, item, q_out):
img = cv2.resize(img, newsize)
try:
- s = mx.recordio.pack_img(header, img, quality=args.quality, img_fmt=args.encoding)
+ if not args.pb_format:
+ s = mx.recordio.pack_img(header, img, quality=args.quality,
+ img_fmt=args.encoding)
+ else:
+ import recordio_pb_pack as rec_pb
+ data = rec_pb.pack_img_data(id=item[0], img=img, quality=args.quality,
+ img_fmt=args.encoding)
+ s = rec_pb.pack(header, data, labels)
q_out.put((i, s, item))
except Exception as e:
traceback.print_exc()
@@ -253,6 +270,8 @@ def parse_args():
help='specify the encoding of the images.')
rgroup.add_argument('--pack-label', type=bool, default=False,
help='Whether to also pack multi dimensional label in the record file')
+ rgroup.add_argument('--pb-format', type=bool, default=False,
+ help='Whether to pack data in pb format')
args = parser.parse_args()
args.prefix = os.path.abspath(args.prefix)
args.root = os.path.abspath(args.root)
diff --git a/tools/recordio_pb_pack.py b/tools/recordio_pb_pack.py
new file mode 100644
index 0000000000..2d459cb894
--- /dev/null
+++ b/tools/recordio_pb_pack.py
@@ -0,0 +1,236 @@
+# -*- coding: utf-8 -*-
+import mxnet.ndarray as nd
+from mxnet.base import numeric_types
+from mxnet.recordio import RecordIOType
+from recordio_pb2 import RecordHead, RecordData, RecordUnit, ExtraData
+try:
+ import cv2
+except ImportError:
+ cv2 = None
+
+def pack_head(id, reserve=0, version=1.0):
+ """ pack a head info into RecordHead format
+
+ Parameters
+ ----------
+ id : int
+ record id.
+ reserve : int
+ record reserve info
+ version : float
+ version of data structure
+
+ Returns
+ -------
+ RecordHead
+ The head format in recordio.proto.
+ """
+
+ head = RecordHead()
+ head.id = id
+ head.reserve = reserve
+ head.version = version
+ return head
+
+def pack_ndarray_data(id, data):
+ """ pack a NDArray into RecordData
+
+ Parameters
+ ----------
+ id : int
+ record id.
+ data : NDArray
+ ndarray data.
+
+ Returns
+ -------
+ RecordData
+ The data format in recordio.proto.
+ """
+ rec_data = RecordData()
+ rec_data.id = id
+ rec_data.type = RecordIOType.NDARRAY
+ rec_data.value = nd.save_to_str(data)
+ return rec_data
+
+def pack_string_data(id, data, is_binary=True):
+ """ pack a string data info into RecordData
+
+ Parameters
+ ----------
+ id : int
+ record id.
+ data : str
+ string data.
+ is_binary : bool
+ whether save the data as binary,
+ NDArray format data will read in RecordIter when using binary,
+
+ Returns
+ -------
+ RecordData
+ The data format in recordio.proto.
+ """
+
+ rec_data = RecordData()
+ rec_data.id = id
+ rec_data.type = RecordIOType.BINARY if is_binary else RecordIOType.STRING
+ rec_data.value = data
+ return rec_data
+
+def pack_ndarray_extra(key, data):
+ """ pack a NDArray into RecordData
+
+ Parameters
+ ----------
+ id : int
+ record id.
+ data : NDArray
+ ndarray data.
+
+ Returns
+ -------
+ RecordData
+ The data format in recordio.proto.
+ """
+ rec_data = ExtraData()
+ rec_data.key = key
+ rec_data.type = RecordIOType.NDARRAY
+ rec_data.value = nd.save_to_str(data)
+ return rec_data
+
+def pack_string_extra(key, data, is_binary=True):
+ """ pack a string extra data info into RecordData
+
+ Parameters
+ ----------
+ id : int
+ record id.
+ data : str
+ string data.
+ is_binary : bool
+ whether save the data as binary,
+ NDArray format data will read in RecordIter when using binary,
+
+ Returns
+ -------
+ RecordData
+ The data format in recordio.proto.
+ """
+
+
+ rec_data = ExtraData()
+ rec_data.key = key
+ rec_data.type = RecordIOType.BINARY if is_binary else RecordIOType.STRING
+ rec_data.value = data
+ return rec_data
+
+def pack(header, data, label, extra=[]):
+ """ pack a data into string for MXRecordIO
+
+ Parameters
+ ----------
+ header : RecordHead
+ Header of the image record.
+ data : RecordData or list of RecordData
+ record data list
+ label : float or list of float
+ label list
+ extra : ExtraData or list of ExtraData
+
+ Returns
+ -------
+ s : str
+ The packed string.
+
+ Examples
+ --------
+ >>> label = 4 # label can also be a 1-D array, for example: label = [1,2,3]
+ >>> id = 2574
+ >>> header = recordio_pb_pack.pack_head(id)
+ >>> with open(path, 'r') as file:
+ ... s = file.read()
+ >>> packed_s = recordio_pb_pack.pack(header, s, label)
+ """
+ if isinstance(data, list):
+ for d in data:
+ if not isinstance(d, RecordData):
+ raise TypeError('type %s not supported' % str(type(d)))
+ elif isinstance(data, RecordData):
+ data = [data]
+ else:
+ raise TypeError('type %s not supported' % str(type(data)))
+
+ if isinstance(label, list):
+ for l in label:
+ if not isinstance(l, numeric_types):
+ raise TypeError('type %s not supported' % str(type(l)))
+ elif isinstance(label, numeric_types):
+ label = [label]
+ else:
+ raise TypeError('type %s not supported' % str(type(label)))
+
+ if isinstance(extra, list):
+ for e in extra:
+ if not isinstance(e, ExtraData):
+ raise TypeError('type %s not supported' % str(type(e)))
+ elif isinstance(extra, ExtraData):
+ extra = [extra]
+ else:
+ raise TypeError('type %s not supported' % str(type(extra)))
+
+ rec_unit = RecordUnit()
+ rec_unit.head.CopyFrom(header)
+
+ # body
+ rec_body = rec_unit.body
+ for d in data:
+ rec_data = rec_body.data.add()
+ rec_data.CopyFrom(d)
+ for l in label:
+ rec_body.label.append(l)
+ for e in extra:
+ ed = rec_body.extra.add()
+ ed.CopyFrom(e)
+ return rec_unit.SerializeToString()
+
+def pack_img_data(id, img, quality=95, img_fmt='.jpg'):
+ """Pack an image into string.
+
+ Parameters
+ ----------
+ id : uint64
+ id of the image record.
+ img : numpy.ndarray
+ Image to be packed.
+ quality : int
+ Quality for JPEG encoding in range 1-100, or compression for PNG encoding in range 1-9.
+ img_fmt : str
+ Encoding of the image (.jpg for JPEG, .png for PNG).
+
+ Returns
+ -------
+ s : str
+ The packed string.
+
+ Examples
+ --------
+ >>> id = 2574
+ >>> label = 4 # label can also be a 1-D array, for example: label = [1,2,3]
+ >>> header = recordio_pb_pack.pack_head(id)
+ >>> img = cv2.imread('test.jpg')
+ >>> packed_img = recordio_pb_pack.pack_img_data(id, img)
+ >>> packed_s = recordio_pb_pack.pack(header, packed_img, label)
+ """
+ assert cv2 is not None
+ jpg_formats = ['.JPG', '.JPEG']
+ png_formats = ['.PNG']
+ encode_params = None
+ if img_fmt.upper() in jpg_formats:
+ encode_params = [cv2.IMWRITE_JPEG_QUALITY, quality]
+ elif img_fmt.upper() in png_formats:
+ encode_params = [cv2.IMWRITE_PNG_COMPRESSION, quality]
+
+ ret, buf = cv2.imencode(img_fmt, img, encode_params)
+ assert ret, 'failed to encode image'
+ return pack_string_data(id, buf.tostring())
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services