def get_iterator(data_shape, use_caffe_data):
    """Generate the iterator of mnist dataset"""
    def get_iterator_impl_mnist(args, kv):
        """return train and val iterators for mnist"""
        # download data
        get_mnist_ubyte()
        flat = False if len(data_shape) != 1 else True

        train = mx.io.MNISTIter(
            image="data/train-images-idx3-ubyte",
            label="data/train-labels-idx1-ubyte",
            input_shape=data_shape,
            batch_size=args.batch_size,
            shuffle=True,
            flat=flat,
            num_parts=kv.num_workers,
            part_index=kv.rank)

        val = mx.io.MNISTIter(
            image="data/t10k-images-idx3-ubyte",
            label="data/t10k-labels-idx1-ubyte",
            input_shape=data_shape,
            batch_size=args.batch_size,
            flat=flat,
            num_parts=kv.num_workers,
            part_index=kv.rank)

        return (train, val)

    def get_iterator_impl_caffe(args, kv):
        flat = False if len(data_shape) != 1 else True
        train = mx.io.CaffeDataIter(
            prototxt=
            'layer { \
                name: "mnist" \
                type: "Data" \
                top: "data" \
                top: "label" \
                include { \
                    phase: TRAIN \
                } \
                transform_param { \
                    scale: 0.00390625 \
                } \
                data_param { \
                    source: "mnist_train_lmdb" \
                    batch_size: 64 \
                    backend: LMDB \
                } \
            }',
            flat=flat,
            num_examples=60000
            # float32 is the default, so left out here in order to illustrate
        )

        val = mx.io.CaffeDataIter(
            prototxt=
            'layer { \
                name: "mnist" \
                type: "Data" \
                top: "data" \
                top: "label" \
                include { \
                    phase: TEST \
                } \
                transform_param { \
                    scale: 0.00390625 \
                } \
                data_param { \
                    source: "mnist_test_lmdb" \
                    batch_size: 100 \
                    backend: LMDB \
                } \
            }',
            flat=flat,
            num_examples=10000,
            dtype="float32"  # float32 is the default
        )

        return train, val

    if use_caffe_data:
        return get_iterator_impl_caffe
    else:
        return get_iterator_impl_mnist
Generate the iterator of mnist dataset