| # 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. |
| |
| import os |
| import tarfile |
| import tempfile |
| import unittest |
| import mxnet as mx |
| import numpy as np |
| import random |
| from mxnet import gluon |
| import platform |
| from mxnet.gluon.data import DataLoader |
| import mxnet.ndarray as nd |
| from mxnet import context |
| from mxnet.gluon.data.dataset import Dataset |
| from mxnet.gluon.data.dataset import ArrayDataset |
| import pytest |
| |
| def test_array_dataset(): |
| X = np.random.uniform(size=(10, 20)) |
| Y = np.random.uniform(size=(10,)) |
| dataset = gluon.data.ArrayDataset(X, Y) |
| loader = gluon.data.DataLoader(dataset, 2) |
| for i, (x, y) in enumerate(loader): |
| assert mx.test_utils.almost_equal(x.asnumpy(), X[i*2:(i+1)*2]) |
| assert mx.test_utils.almost_equal(y.asnumpy(), Y[i*2:(i+1)*2]) |
| |
| dataset = gluon.data.ArrayDataset(X) |
| loader = gluon.data.DataLoader(dataset, 2) |
| |
| for i, x in enumerate(loader): |
| assert mx.test_utils.almost_equal(x.asnumpy(), X[i*2:(i+1)*2]) |
| |
| @pytest.fixture(scope="session") |
| def prepare_record(tmpdir_factory): |
| test_images = tmpdir_factory.mktemp("test_images") |
| test_images_tar = test_images.join("test_images.tar.gz") |
| gluon.utils.download("https://repo.mxnet.io/gluon/dataset/test/test_images-9cebe48a.tar.gz", str(test_images_tar)) |
| tarfile.open(test_images_tar).extractall(str(test_images)) |
| imgs = os.listdir(str(test_images.join("test_images"))) |
| record = mx.recordio.MXIndexedRecordIO(str(test_images.join("test.idx")), str(test_images.join("test.rec")), 'w') |
| for i, img in enumerate(imgs): |
| with open(str(test_images.join("test_images").join(img)), 'rb') as f: |
| str_img = f.read() |
| s = mx.recordio.pack((0, i, i, 0), str_img) |
| record.write_idx(i, s) |
| return str(test_images.join('test.rec')) |
| |
| |
| def test_recordimage_dataset(prepare_record): |
| recfile = prepare_record |
| fn = lambda x, y : (x, y) |
| dataset = gluon.data.vision.ImageRecordDataset(recfile).transform(fn) |
| loader = gluon.data.DataLoader(dataset, 1) |
| |
| for i, (x, y) in enumerate(loader): |
| assert x.shape[0] == 1 and x.shape[3] == 3 |
| assert y.asscalar() == i |
| |
| @mx.util.use_np |
| def test_recordimage_dataset_handle(prepare_record): |
| recfile = prepare_record |
| class TmpTransform(mx.gluon.HybridBlock): |
| def forward(self, x): |
| return x |
| |
| fn = TmpTransform() |
| dataset = gluon.data.vision.ImageRecordDataset(recfile).transform_first(fn).__mx_handle__() |
| loader = gluon.data.DataLoader(dataset, 1) |
| |
| for i, (x, y) in enumerate(loader): |
| assert x.shape[0] == 1 and x.shape[3] == 3 |
| assert y.item() == i |
| |
| def _dataset_transform_fn(x, y): |
| """Named transform function since lambda function cannot be pickled.""" |
| return x, y |
| |
| def _dataset_transform_first_fn(x): |
| """Named transform function since lambda function cannot be pickled.""" |
| return x |
| |
| def test_recordimage_dataset_with_data_loader_multiworker(prepare_record): |
| recfile = prepare_record |
| dataset = gluon.data.vision.ImageRecordDataset(recfile) |
| loader = gluon.data.DataLoader(dataset, 1, num_workers=5, try_nopython=False) |
| |
| for i, (x, y) in enumerate(loader): |
| assert x.shape[0] == 1 and x.shape[3] == 3 |
| assert y.asscalar() == i |
| |
| # with transform |
| dataset = gluon.data.vision.ImageRecordDataset(recfile).transform(_dataset_transform_fn) |
| loader = gluon.data.DataLoader(dataset, 1, num_workers=5, try_nopython=None) |
| |
| for i, (x, y) in enumerate(loader): |
| assert x.shape[0] == 1 and x.shape[3] == 3 |
| assert y.asscalar() == i |
| |
| # with transform_first |
| dataset = gluon.data.vision.ImageRecordDataset(recfile).transform_first(_dataset_transform_first_fn) |
| loader = gluon.data.DataLoader(dataset, 1, num_workers=5, try_nopython=None) |
| |
| for i, (x, y) in enumerate(loader): |
| assert x.shape[0] == 1 and x.shape[3] == 3 |
| assert y.asscalar() == i |
| |
| def test_sampler(): |
| seq_sampler = gluon.data.SequentialSampler(10) |
| assert list(seq_sampler) == list(range(10)) |
| rand_sampler = gluon.data.RandomSampler(10) |
| assert sorted(list(rand_sampler)) == list(range(10)) |
| seq_batch_keep = gluon.data.BatchSampler(seq_sampler, 3, 'keep') |
| assert sum(list(seq_batch_keep), []) == list(range(10)) |
| seq_batch_discard = gluon.data.BatchSampler(seq_sampler, 3, 'discard') |
| assert sum(list(seq_batch_discard), []) == list(range(9)) |
| rand_batch_keep = gluon.data.BatchSampler(rand_sampler, 3, 'keep') |
| assert sorted(sum(list(rand_batch_keep), [])) == list(range(10)) |
| |
| def test_datasets(tmpdir): |
| p = tmpdir.mkdir("test_datasets") |
| assert len(gluon.data.vision.MNIST(root=str(p.join('mnist')))) == 60000 |
| assert len(gluon.data.vision.MNIST(root=str(p.join('mnist')), train=False)) == 10000 |
| assert len(gluon.data.vision.FashionMNIST(root=str(p.join('fashion-mnist')))) == 60000 |
| assert len(gluon.data.vision.FashionMNIST(root=str(p.join('fashion-mnist')), train=False)) == 10000 |
| assert len(gluon.data.vision.CIFAR10(root=str(p.join('cifar10')))) == 50000 |
| assert len(gluon.data.vision.CIFAR10(root=str(p.join('cifar10')), train=False)) == 10000 |
| assert len(gluon.data.vision.CIFAR100(root=str(p.join('cifar100')))) == 50000 |
| assert len(gluon.data.vision.CIFAR100(root=str(p.join('cifar100')), fine_label=True)) == 50000 |
| assert len(gluon.data.vision.CIFAR100(root=str(p.join('cifar100')), train=False)) == 10000 |
| |
| def test_datasets_handles(tmpdir): |
| p = tmpdir.mkdir("test_datasets_handles") |
| assert len(gluon.data.vision.MNIST(root=str(p.join('mnist'))).__mx_handle__()) == 60000 |
| assert len(gluon.data.vision.MNIST(root=str(p.join('mnist')), train=False).__mx_handle__()) == 10000 |
| assert len(gluon.data.vision.FashionMNIST(root=str(p.join('fashion-mnist'))).__mx_handle__()) == 60000 |
| assert len(gluon.data.vision.FashionMNIST(root=str(p.join('fashion-mnist')), train=False).__mx_handle__()) == 10000 |
| assert len(gluon.data.vision.CIFAR10(root=str(p.join('cifar10'))).__mx_handle__()) == 50000 |
| assert len(gluon.data.vision.CIFAR10(root=str(p.join('cifar10')), train=False).__mx_handle__()) == 10000 |
| assert len(gluon.data.vision.CIFAR100(root=str(p.join('cifar100'))).__mx_handle__()) == 50000 |
| assert len(gluon.data.vision.CIFAR100(root=str(p.join('cifar100')), fine_label=True).__mx_handle__()) == 50000 |
| assert len(gluon.data.vision.CIFAR100(root=str(p.join('cifar100')), train=False).__mx_handle__()) == 10000 |
| |
| def test_image_folder_dataset(prepare_record): |
| dataset = gluon.data.vision.ImageFolderDataset(os.path.dirname(prepare_record)) |
| assert dataset.synsets == ['test_images'] |
| assert len(dataset.items) == 16 |
| |
| def test_image_folder_dataset_handle(prepare_record): |
| dataset = gluon.data.vision.ImageFolderDataset(os.path.dirname(prepare_record)) |
| hd = dataset.__mx_handle__() |
| assert len(hd) == 16 |
| assert (hd[1][0] == dataset[1][0]).asnumpy().all() |
| assert hd[5][1] == dataset[5][1] |
| |
| def test_image_list_dataset(prepare_record): |
| root = os.path.join(os.path.dirname(prepare_record), 'test_images') |
| imlist = os.listdir(root) |
| imglist = [(0, path) for i, path in enumerate(imlist)] |
| dataset = gluon.data.vision.ImageListDataset(root=root, imglist=imglist) |
| assert len(dataset) == 16, len(dataset) |
| img, label = dataset[0] |
| assert len(img.shape) == 3 |
| assert label == 0 |
| |
| # save to file as *.lst |
| imglist = ['\t'.join((str(i), '0', path)) for i, path in enumerate(imlist)] |
| with tempfile.NamedTemporaryFile('wt', delete=False) as fp: |
| for line in imglist: |
| fp.write(line + '\n') |
| fp.close() |
| |
| dataset = gluon.data.vision.ImageListDataset(root=root, imglist=fp.name) |
| assert len(dataset) == 16, len(dataset) |
| img, label = dataset[0] |
| assert len(img.shape) == 3 |
| assert label == 0 |
| |
| def test_image_list_dataset_handle(prepare_record): |
| root = os.path.join(os.path.dirname(prepare_record), 'test_images') |
| imlist = os.listdir(root) |
| imglist = [(0, path) for i, path in enumerate(imlist)] |
| dataset = gluon.data.vision.ImageListDataset(root=root, imglist=imglist).__mx_handle__() |
| assert len(dataset) == 16, len(dataset) |
| img, label = dataset[0] |
| assert len(img.shape) == 3 |
| assert label == 0 |
| |
| # save to file as *.lst |
| imglist = ['\t'.join((str(i), '0', path)) for i, path in enumerate(imlist)] |
| with tempfile.NamedTemporaryFile('wt', delete=False) as fp: |
| for line in imglist: |
| fp.write(line + '\n') |
| fp.close() |
| |
| dataset = gluon.data.vision.ImageListDataset(root=root, imglist=fp.name).__mx_handle__() |
| assert len(dataset) == 16 |
| img, label = dataset[0] |
| assert len(img.shape) == 3 |
| assert label == 0 |
| |
| @pytest.mark.garbage_expected |
| def test_list_dataset(): |
| for num_worker in range(0, 3): |
| data = mx.gluon.data.DataLoader([([1,2], 0), ([3, 4], 1)], batch_size=1, num_workers=num_worker) |
| for _ in data: |
| pass |
| |
| |
| class _Dataset(gluon.data.Dataset): |
| def __len__(self): |
| return 100 |
| def __getitem__(self, key): |
| return mx.nd.full((10,), key) |
| |
| @pytest.mark.garbage_expected |
| def test_multi_worker(): |
| data = _Dataset() |
| for thread_pool in [True, False]: |
| loader = gluon.data.DataLoader(data, batch_size=1, num_workers=5, thread_pool=thread_pool) |
| for i, batch in enumerate(loader): |
| assert (batch.asnumpy() == i).all() |
| |
| |
| def test_multi_worker_shape(): |
| for thread_pool in [True, False]: |
| batch_size = 1024 |
| shape = (batch_size+1, 11, 12) |
| |
| data = ArrayDataset(np.ones(shape)) |
| loader = gluon.data.DataLoader( |
| data, batch_size=batch_size, num_workers=5, last_batch='keep', thread_pool=thread_pool) |
| for batch in loader: |
| if shape[0] > batch_size: |
| assert batch.shape == (batch_size, shape[1], shape[2]) |
| shape = (shape[0] - batch_size, shape[1], shape[2]) |
| else: |
| assert batch.shape == shape |
| |
| class _Dummy(Dataset): |
| """Dummy dataset for randomized shape arrays.""" |
| def __init__(self, random_shape): |
| self.random_shape = random_shape |
| |
| def __getitem__(self, idx): |
| key = idx |
| if self.random_shape: |
| out = np.random.uniform(size=(random.randint(1000, 1100), 40)) |
| labels = np.random.uniform(size=(random.randint(10, 15))) |
| else: |
| out = np.random.uniform(size=(1000, 40)) |
| labels = np.random.uniform(size=(10)) |
| return key, out, labels |
| |
| def __len__(self): |
| return 50 |
| |
| def _batchify_list(data): |
| """ |
| return list of ndarray without stack/concat/pad |
| """ |
| if isinstance(data, (tuple, list)): |
| return list(data) |
| if isinstance(data, mx.nd.NDArray): |
| return [data] |
| return data |
| |
| def _batchify(data): |
| """ |
| Collate data into batch. Use shared memory for stacking. |
| :param data: a list of array, with layout of 'NTC'. |
| :return either x and x's unpadded lengths, or x, x's unpadded lengths, y and y's unpadded lengths |
| if labels are not supplied. |
| """ |
| |
| # input layout is NTC |
| keys, inputs, labels = [item[0] for item in data], [item[1] for item in data], \ |
| [item[2] for item in data] |
| |
| if len(data) > 1: |
| max_data_len = max([seq.shape[0] for seq in inputs]) |
| max_labels_len = 0 if not labels else max([seq.shape[0] for seq in labels]) |
| else: |
| max_data_len = inputs[0].shape[0] |
| max_labels_len = 0 if not labels else labels[0].shape[0] |
| |
| x_lens = [item.shape[0] for item in inputs] |
| y_lens = [item.shape[0] for item in labels] |
| |
| for i, seq in enumerate(inputs): |
| pad_len = max_data_len - seq.shape[0] |
| inputs[i] = np.pad(seq, ((0, pad_len), (0, 0)), 'constant', constant_values=0) |
| labels[i] = np.pad(labels[i], (0, max_labels_len - labels[i].shape[0]), |
| 'constant', constant_values=-1) |
| |
| inputs = np.asarray(inputs, dtype=np.float32) |
| if labels is not None: |
| labels = np.asarray(labels, dtype=np.float32) |
| inputs = inputs.transpose((1, 0, 2)) |
| labels = labels.transpose((1, 0)) |
| |
| return (nd.array(inputs, dtype=inputs.dtype, ctx=context.Context('cpu_shared', 0)), |
| nd.array(x_lens, ctx=context.Context('cpu_shared', 0))) \ |
| if labels is None else ( |
| nd.array(inputs, dtype=inputs.dtype, ctx=context.Context('cpu_shared', 0)), |
| nd.array(x_lens, ctx=context.Context('cpu_shared', 0)), |
| nd.array(labels, dtype=labels.dtype, ctx=context.Context('cpu_shared', 0)), |
| nd.array(y_lens, ctx=context.Context('cpu_shared', 0))) |
| |
| def test_multi_worker_forked_data_loader(): |
| data = _Dummy(False) |
| loader = DataLoader(data, batch_size=40, batchify_fn=_batchify, num_workers=2) |
| for _ in range(1): |
| for _ in loader: |
| pass |
| |
| data = _Dummy(True) |
| loader = DataLoader(data, batch_size=40, batchify_fn=_batchify_list, num_workers=2) |
| for _ in range(1): |
| for _ in loader: |
| pass |
| |
| def test_multi_worker_dataloader_release_pool(): |
| # will trigger too many open file if pool is not released properly |
| if os.name == 'nt': |
| print('Skip for windows since spawn on windows is too expensive.') |
| return |
| |
| for _ in range(10): |
| A = np.random.rand(999, 2000) |
| D = mx.gluon.data.DataLoader(A, batch_size=8, num_workers=8) |
| the_iter = iter(D) |
| next(the_iter) |
| del the_iter |
| del D |
| |
| def test_dataloader_context(): |
| X = np.random.uniform(size=(10, 20)) |
| dataset = gluon.data.ArrayDataset(X) |
| default_dev_id = 0 |
| custom_dev_id = 1 |
| |
| # use non-pinned memory |
| loader1 = gluon.data.DataLoader(dataset, 8) |
| for _, x in enumerate(loader1): |
| assert x.context == context.cpu(default_dev_id) |
| |
| # use pinned memory with default device id |
| loader2 = gluon.data.DataLoader(dataset, 8, pin_memory=True) |
| for _, x in enumerate(loader2): |
| assert x.context == context.cpu_pinned(default_dev_id) |
| |
| if mx.device.num_gpus() <= 1: |
| print('Bypassing custom_dev_id pinned mem test on system with < 2 gpus.') |
| else: |
| # use pinned memory with custom device id |
| loader3 = gluon.data.DataLoader(dataset, 8, pin_memory=True, |
| pin_device_id=custom_dev_id) |
| for _, x in enumerate(loader3): |
| assert x.context == context.cpu_pinned(custom_dev_id) |
| |
| def batchify(a): |
| return a |
| |
| def test_dataset_filter(): |
| length = 100 |
| a = mx.gluon.data.SimpleDataset([i for i in range(length)]) |
| a_filtered = a.filter(lambda x: x % 10 == 0) |
| assert(len(a_filtered) == 10) |
| for sample in a_filtered: |
| assert sample % 10 == 0 |
| a_xform_filtered = a.transform(lambda x: x + 1).filter(lambda x: x % 10 == 0) |
| assert(len(a_xform_filtered) == 10) |
| # the filtered data is already transformed |
| for sample in a_xform_filtered: |
| assert sample % 10 == 0 |
| |
| def test_dataset_filter_handle(): |
| length = 100 |
| a = mx.gluon.data.SimpleDataset(np.arange(length)) |
| a_filtered = a.filter(lambda x: x % 10 == 0).__mx_handle__() |
| assert(len(a_filtered) == 10) |
| for sample in a_filtered: |
| assert sample % 10 == 0 |
| a_xform_filtered = a.transform(lambda x: x + 1).filter(lambda x: x % 10 == 0) |
| assert(len(a_xform_filtered) == 10) |
| # the filtered data is already transformed |
| for sample in a_xform_filtered: |
| assert sample % 10 == 0 |
| |
| def test_dataset_shard(): |
| length = 9 |
| a = mx.gluon.data.SimpleDataset([i for i in range(length)]) |
| shard_0 = a.shard(4, 0) |
| shard_1 = a.shard(4, 1) |
| shard_2 = a.shard(4, 2) |
| shard_3 = a.shard(4, 3) |
| assert len(shard_0) + len(shard_1) + len(shard_2) + len(shard_3) == length |
| assert len(shard_0) == 3 |
| assert len(shard_1) == 2 |
| assert len(shard_2) == 2 |
| assert len(shard_3) == 2 |
| total = 0 |
| for shard in [shard_0, shard_1, shard_2, shard_3]: |
| for sample in shard: |
| total += sample |
| assert total == sum(a) |
| |
| def test_dataset_shard_handle(): |
| length = 9 |
| a = mx.gluon.data.SimpleDataset(np.arange(length)) |
| shard_0 = a.shard(4, 0).__mx_handle__() |
| shard_1 = a.shard(4, 1).__mx_handle__() |
| shard_2 = a.shard(4, 2).__mx_handle__() |
| shard_3 = a.shard(4, 3).__mx_handle__() |
| assert len(shard_0) + len(shard_1) + len(shard_2) + len(shard_3) == length |
| assert len(shard_0) == 3 |
| assert len(shard_1) == 2 |
| assert len(shard_2) == 2 |
| assert len(shard_3) == 2 |
| total = 0 |
| for shard in [shard_0, shard_1, shard_2, shard_3]: |
| for sample in shard: |
| total += sample |
| assert total == sum(a) |
| |
| def test_dataset_take(): |
| length = 100 |
| a = mx.gluon.data.SimpleDataset([i for i in range(length)]) |
| a_take_full = a.take(1000) |
| assert len(a_take_full) == length |
| a_take_full = a.take(None) |
| assert len(a_take_full) == length |
| count = 10 |
| a_take_10 = a.take(count) |
| assert len(a_take_10) == count |
| expected_total = sum([i for i in range(count)]) |
| total = 0 |
| for sample in a_take_10: |
| assert sample < count |
| total += sample |
| assert total == expected_total |
| |
| a_xform_take_10 = a.transform(lambda x: x * 10).take(count) |
| assert len(a_xform_take_10) == count |
| expected_total = sum([i * 10 for i in range(count)]) |
| total = 0 |
| for sample in a_xform_take_10: |
| assert sample < count * 10 |
| total += sample |
| assert total == expected_total |
| |
| def test_dataset_take_handle(): |
| length = 100 |
| a = mx.gluon.data.SimpleDataset(np.arange(length)) |
| a_take_full = a.take(1000).__mx_handle__() |
| assert len(a_take_full) == length |
| a_take_full = a.take(None).__mx_handle__() |
| assert len(a_take_full) == length |
| count = 10 |
| a_take_10 = a.take(count).__mx_handle__() |
| assert len(a_take_10) == count |
| expected_total = sum([i for i in range(count)]) |
| total = 0 |
| for sample in a_take_10: |
| assert sample < count |
| total += sample |
| assert total == expected_total |
| |
| a_xform_take_10 = a.take(count).__mx_handle__() |
| assert len(a_xform_take_10) == count |
| expected_total = sum([i for i in range(count)]) |
| total = 0 |
| for sample in a_xform_take_10: |
| assert sample < count |
| total += sample |
| assert total == expected_total |
| |
| @pytest.mark.garbage_expected |
| def test_dataloader_scope(): |
| """ |
| Bug: Gluon DataLoader terminates the process pool early while |
| _MultiWorkerIter is operating on the pool. |
| |
| Tests that DataLoader is not garbage collected while the iterator is |
| in use. |
| """ |
| args = {'num_workers': 1, 'batch_size': 2} |
| dataset = nd.ones(5) |
| iterator = iter(DataLoader( |
| dataset, |
| batchify_fn=batchify, |
| **args |
| ) |
| ) |
| |
| item = next(iterator) |
| |
| assert item is not None |
| |
| def test_mx_datasets_handle(): |
| # _DownloadedDataset |
| mnist = mx.gluon.data.vision.MNIST(train=False).__mx_handle__() |
| assert len(mnist) == 10000 |
| cifar10 = mx.gluon.data.vision.CIFAR10(train=False).__mx_handle__() |
| assert len(cifar10) == 10000 |
| |
| # _SampledDataset |
| s_mnist = mnist.take(100).__mx_handle__() |
| assert len(s_mnist) == 100 |
| assert np.all(s_mnist[0][0].asnumpy() == mnist[0][0].asnumpy()) |
| assert s_mnist[0][1] == mnist[0][1] |
| |
| # ArrayDataset |
| mc = mx.gluon.data.ArrayDataset(mnist.take(100), cifar10.take(100)).__mx_handle__() |
| assert len(mc) == 100 |
| assert len(mc[0]) == 4 # two from mnist, two from cifar10 |
| assert mc[0][1] == mnist[0][1] |
| assert mc[0][3] == cifar10[0][1] |
| |
| def test_mx_data_loader(): |
| from mxnet.gluon.data.dataloader import DataLoader |
| |
| dataset = mx.gluon.data.vision.MNIST(train=False) |
| dl = DataLoader(num_workers=0, dataset=dataset, batch_size=32) |
| for _ in dl: |
| pass |
| |
| @mx.util.use_np |
| def test_mx_data_loader_nopython(): |
| from mxnet.gluon.data.dataloader import DataLoader |
| from mxnet.gluon.data.vision.transforms import ToTensor |
| dataset = mx.gluon.data.vision.MNIST(train=False) |
| dl1 = DataLoader(dataset=dataset.transform_first(ToTensor()), batch_size=32, try_nopython=True, shuffle=False) |
| dl2 = DataLoader(dataset=dataset.transform_first(ToTensor()), batch_size=32, try_nopython=False, shuffle=False) |
| assert len(dl1) == len(dl2) |
| assert np.all(next(iter(dl1))[1].asnumpy() == next(iter(dl2))[1].asnumpy()) |
| for _ in dl1: |
| pass |
| |
| def test_batchify_stack(): |
| a = np.array([[1, 2, 3, 4], [5, 6, 7, 8]]) |
| b = np.array([[5, 6, 7, 8], [1, 2, 3, 4]]) |
| bf = mx.gluon.data.batchify.Stack() |
| bf_handle = bf.__mx_handle__() |
| c = bf([a, b]) |
| d = bf_handle([a, b]) |
| assert c.shape == d.shape |
| assert mx.test_utils.almost_equal(c.asnumpy(), d.asnumpy()) |
| assert mx.test_utils.almost_equal(c.asnumpy(), np.stack((a, b))) |
| |
| def test_batchify_pad(): |
| a = np.array([[1, 2, 3, 4], [11, 12, 13, 14]]) |
| b = np.array([[4, 5, 6]]) |
| c = np.array([[9, 10]]) |
| bf = mx.gluon.data.batchify.Pad(val=-1) |
| bf_handle = bf.__mx_handle__() |
| d = bf([a, b, c]) |
| e = bf_handle([a, b, c]) |
| assert d.shape == e.shape |
| assert mx.test_utils.almost_equal(d.asnumpy(), e.asnumpy()) |
| expected = np.array([[[ 1., 2., 3., 4.], [11., 12., 13., 14.]], |
| [[ 4., 5., 6., -1.], [-1., -1., -1., -1.]], |
| [[ 9., 10., -1., -1.], [-1., -1., -1., -1.]]]) |
| assert mx.test_utils.almost_equal(d.asnumpy(), expected) |
| |
| def test_batchify_group(): |
| a = [np.array([[1, 2, 3, 4], [5, 6, 7, 8]]), np.array([[1, 2, 3, 4], [11, 12, 13, 14]])] |
| b = [np.array([[1, 2, 3, 4], [5, 6, 7, 8]]), np.array([[4, 5, 6]])] |
| c = [np.array([[1, 2, 3, 4], [5, 6, 7, 8]]), np.array([[9, 10]])] |
| bf = mx.gluon.data.batchify.Group(mx.gluon.data.batchify.Stack(), mx.gluon.data.batchify.Pad(val=-1)) |
| bf_handle = bf.__mx_handle__() |
| d = bf([a, b, c]) |
| e = bf_handle([a, b, c]) |
| assert d[0].shape == e[0].shape |
| assert d[1].shape == e[1].shape |
| print(d[0].asnumpy(), ',', e[0].asnumpy(), ',', e[1].asnumpy()) |
| assert mx.test_utils.almost_equal(d[0].asnumpy(), e[0].asnumpy()) |
| assert mx.test_utils.almost_equal(d[1].asnumpy(), e[1].asnumpy()) |
| assert mx.test_utils.almost_equal(d[0].asnumpy(), np.stack((a[0], b[0], c[0]))) |
| expected = np.array([[[ 1., 2., 3., 4.], [11., 12., 13., 14.]], |
| [[ 4., 5., 6., -1.], [-1., -1., -1., -1.]], |
| [[ 9., 10., -1., -1.], [-1., -1., -1., -1.]]]) |
| assert mx.test_utils.almost_equal(d[1].asnumpy(), expected) |
| |
| def test_sampler(): |
| interval_sampler = mx.gluon.data.IntervalSampler(10, 3) |
| assert sorted(list(interval_sampler)) == list(range(10)) |
| interval_sampler = mx.gluon.data.IntervalSampler(10, 3, rollover=False) |
| assert list(interval_sampler) == [0, 3, 6, 9] |
