| import sys |
| import os |
| # code to automatically download dataset |
| curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__))) |
| sys.path.append(os.path.join(curr_path, "../../tests/python/common")) |
| import get_data |
| import mxnet as mx |
| |
| def get_iterator(data_shape, use_caffe_data): |
| def get_iterator_impl_mnist(args, kv): |
| """return train and val iterators for mnist""" |
| # download data |
| get_data.GetMNIST_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 |