def test_mxnet_iterator_model_fit():
from nvidia.dali.plugin.mxnet import DALIGenericIterator as MXNetIterator
import mxnet as mx
num_gpus = 1
batch_size = 1
class RN50Pipeline(Pipeline):
def __init__(self, batch_size, num_threads, device_id, num_gpus, data_paths):
super(RN50Pipeline, self).__init__(batch_size, num_threads, device_id,)
self.input = ops.FileReader(file_root = data_paths, shard_id = device_id, num_shards = num_gpus)
self.decode_gpu = ops.ImageDecoder(device = "cpu", output_type = types.RGB)
self.res = ops.RandomResizedCrop(device="cpu", size =(224,224))
self.cmnp = ops.CropMirrorNormalize(device="cpu",
output_dtype=types.FLOAT,
output_layout=types.NCHW,
crop=(224, 224),
image_type=types.RGB,
mean=[0.485 * 255,0.456 * 255,0.406 * 255],
std=[0.229 * 255,0.224 * 255,0.225 * 255])
self.coin = ops.CoinFlip(probability=0.5)
def define_graph(self):
rng = self.coin()
jpegs, labels = self.input(name="Reader")
images = self.decode_gpu(jpegs)
images = self.res(images)
output = self.cmnp(images, mirror=rng)
return labels
pipes, _ = create_pipeline(lambda gpu: RN50Pipeline(batch_size=batch_size, num_threads=4, device_id=gpu, num_gpus=num_gpus,
data_paths=image_data_set), batch_size, num_gpus)
pipe = pipes[0]
class MXNetIteratorWrapper(MXNetIterator):
def __init__(self, iter):
self.iter = iter
def __getattr__(self, attr):
return getattr(self.iter, attr)
def __next__(self):
ret = self.iter.__next__()[0]
return ret
dali_train_iter = MXNetIterator(pipe, [("labels", MXNetIterator.LABEL_TAG)],
size=pipe.epoch_size("Reader"))
data = mx.symbol.Variable('labels')
# create a dummy model
_ = mx.model.FeedForward.create(data,
X=MXNetIteratorWrapper(dali_train_iter),
num_epoch=1,
learning_rate=0.01)
def test_mxnet_pipeline_dynamic_shape():
from nvidia.dali.plugin.mxnet import DALIGenericIterator as MXNetIterator
root, annotations = data_paths()
pipeline = DetectionPipeline(BATCH_SIZE, 0, root, annotations)
train_loader = MXNetIterator([pipeline],
[('data', MXNetIterator.DATA_TAG),
('bboxes', MXNetIterator.LABEL_TAG),
('label', MXNetIterator.LABEL_TAG)],
EPOCH_SIZE,
auto_reset=False,
dynamic_shape=True)
for data in train_loader:
assert data is not None
def test_stop_iteration():
for fw_iter, iter_name in [(lambda pipe, size, auto_reset: PyTorchIterator(
pipe, output_map=["data"], size=size, auto_reset=auto_reset),
"PyTorchIterator"),
(lambda pipe, size, auto_reset: MXNetIterator(
pipe, [("data", MXNetIterator.DATA_TAG)],
size=size,
auto_reset=auto_reset), "MXNetIterator")]:
for epochs in [1, 3, 6]:
for iter_num in [2, 5, 9]:
for batch_size in [1, 10, 100]:
for auto_reset in [True, False]:
yield check_stop_iter, fw_iter, iter_name, batch_size, epochs, iter_num, auto_reset
def test_mxnet_iterator_empty_array():
from nvidia.dali.plugin.mxnet import DALIGenericIterator as MXNetIterator
import mxnet as mx
batch_size = 4
size = 5
all_np_types = [np.bool_, np.int_, np.intc, np.intp, np.int8, np.int16, np.int32, np.int64,
np.uint8, np.uint16, np.uint32, np.uint64, np.float_, np.float32, np.float16,
np.short, np.long, np.longlong, np.ushort, np.ulonglong]
np_types = []
# store in np_types only types supported by MXNet
for t in all_np_types:
try:
mx.nd.zeros([2, 2, 2], ctx=None, dtype=t)
np_types.append(t)
except mx.base.MXNetError:
pass
test_data_shape = [1, 3, 0, 4]
def get_data():
# create batch of [type_a, type_a, type_b, type_b, ...]
out = [[np.empty(test_data_shape, dtype = t)] * batch_size for t in np_types]
out = [val for pair in zip(out, out) for val in pair]
return out
pipe = Pipeline(batch_size=batch_size, num_threads=3, device_id=0)
outs = fn.external_source(source = get_data, num_outputs = len(np_types) * 2)
pipe.set_outputs(*outs)
pipe.build()
# create map of [(data, type_a), (label, type_a), ...]
data_map = [('data_{}'.format(i), MXNetIterator.DATA_TAG) for i, t in enumerate(np_types)]
label_map = [('label_{}'.format(i), MXNetIterator.LABEL_TAG) for i, t in enumerate(np_types)]
out_map = [val for pair in zip(data_map, label_map) for val in pair]
iterator = MXNetIterator(
pipe,
output_map=out_map,
size=size,
dynamic_shape=True)
for batch in iterator:
for d, t in zip(batch[0].data, np_types):
shape = d.asnumpy().shape
assert shape[0] == batch_size
print(shape)
assert np.array_equal(shape[1:], test_data_shape)
assert d.asnumpy().dtype == t
def test_mxnet_iterator_last_batch_no_pad_last_batch():
num_gpus = 1
batch_size = 100
pipes, data_size = create_pipeline(lambda gpu: COCOReaderPipeline(batch_size=batch_size, num_threads=4, shard_id=gpu, num_gpus=num_gpus,
data_paths=data_sets[0], random_shuffle=True, stick_to_shard=False,
shuffle_after_epoch=False, pad_last_batch=False), batch_size, num_gpus)
dali_train_iter = MXNetIterator(pipes, [("ids", MXNetIterator.DATA_TAG)],
size=pipes[0].epoch_size("Reader"), fill_last_batch=True)
img_ids_list, img_ids_list_set, mirrored_data, _, _ = \
gather_ids(dali_train_iter, lambda x: x.data[0].squeeze().asnumpy(), lambda x: x.pad, data_size)
assert len(img_ids_list) > data_size
assert len(img_ids_list_set) == data_size
assert len(set(mirrored_data)) != 1
def test_mxnet_iterator_not_fill_last_batch_pad_last_batch():
from nvidia.dali.plugin.mxnet import DALIGenericIterator as MXNetIterator
num_gpus = 1
batch_size = 100
pipes, data_size = create_pipeline(
lambda gpu: COCOReaderPipeline(batch_size=batch_size,
num_threads=4,
shard_id=gpu,
num_gpus=num_gpus,
data_paths=data_sets[0],
random_shuffle=True,
stick_to_shard=False,
shuffle_after_epoch=False,
pad_last_batch=True), batch_size,
num_gpus)
dali_train_iter = MXNetIterator(pipes, [("ids", MXNetIterator.DATA_TAG)],
size=pipes[0].epoch_size("Reader"),
fill_last_batch=False)
img_ids_list, img_ids_list_set, mirrored_data, pad, remainder = \
gather_ids(dali_train_iter, lambda x: x.data[0].squeeze().asnumpy(), lambda x: x.pad, data_size)
assert pad == remainder
assert len(img_ids_list) - pad == data_size
assert len(img_ids_list_set) == data_size
assert len(set(mirrored_data)) == 1
dali_train_iter.reset()
next_img_ids_list, next_img_ids_list_set, next_mirrored_data, pad, remainder = \
gather_ids(dali_train_iter, lambda x: x.data[0].squeeze().asnumpy(), lambda x: x.pad, data_size)
assert pad == remainder
assert len(next_img_ids_list) - pad == data_size
assert len(next_img_ids_list_set) == data_size
assert len(set(next_mirrored_data)) == 1
def test_stop_iteration_mxnet():
from nvidia.dali.plugin.mxnet import DALIGenericIterator as MXNetIterator
test_stop_iteration(lambda pipe, size, auto_reset : MXNetIterator(pipe, [("data", MXNetIterator.DATA_TAG)], size=size, auto_reset=auto_reset), "MXNetIterator")
def test_stop_iteration_mxnet_fail_single():
from nvidia.dali.plugin.mxnet import DALIGenericIterator as MXNetIterator
fw_iter = lambda pipe, size, auto_reset : MXNetIterator(pipe, [("data", MXNetIterator.DATA_TAG)], size=size, auto_reset=auto_reset)
check_stop_iter_fail_single(fw_iter)
def test_stop_iteration_mxnet():
from nvidia.dali.plugin.mxnet import DALIGenericIterator as MXNetIterator
fw_iter = lambda pipe, size, auto_reset : MXNetIterator(pipe, [("data", MXNetIterator.DATA_TAG)], size=size, auto_reset=auto_reset)
iter_name = "MXNetIterator"
for batch_size, epochs, iter_num, auto_reset, infinite in stop_teration_case_generator():
yield check_stop_iter, fw_iter, iter_name, batch_size, epochs, iter_num, auto_reset, infinite