Esempio n. 1
0
def test_sharded_parallel_sample_iter_num_batches():
    num_shards = 2
    batch_size = 2
    num_batches_per_bucket = 10
    buckets = data_io_pt.define_parallel_buckets(100, 100, 10, True, 1.0)
    bucket_counts = [batch_size * num_batches_per_bucket for _ in buckets]
    num_batches_per_shard = num_batches_per_bucket * len(buckets)
    num_batches = num_shards * num_batches_per_shard
    bucket_batch_sizes = data_io_pt.define_bucket_batch_sizes(
        buckets,
        batch_size,
        batch_type=C.BATCH_TYPE_SENTENCE,
        data_target_average_len=[None] * len(buckets))

    dataset1 = data_io_pt.ParallelDataSet(*_get_random_bucketed_data(
        buckets, min_count=0, max_count=5, bucket_counts=bucket_counts))
    dataset2 = data_io_pt.ParallelDataSet(*_get_random_bucketed_data(
        buckets, min_count=0, max_count=5, bucket_counts=bucket_counts))
    with TemporaryDirectory() as work_dir:
        shard1_fname = os.path.join(work_dir, 'shard1')
        shard2_fname = os.path.join(work_dir, 'shard2')
        dataset1.save(shard1_fname)
        dataset2.save(shard2_fname)
        shard_fnames = [shard1_fname, shard2_fname]

        it = data_io_pt.ShardedParallelSampleIter(shard_fnames, buckets,
                                                  batch_size,
                                                  bucket_batch_sizes)

        num_batches_seen = 0
        while it.iter_next():
            it.next()
            num_batches_seen += 1
        assert num_batches_seen == num_batches
Esempio n. 2
0
def test_parallel_sample_iter():
    batch_size = 2
    buckets = data_io_pt.define_parallel_buckets(100, 100, 10, True, 1.0)
    # The first bucket is going to be empty:
    bucket_counts = [0] + [None] * (len(buckets) - 1)
    bucket_batch_sizes = data_io_pt.define_bucket_batch_sizes(
        buckets,
        batch_size,
        batch_type=C.BATCH_TYPE_SENTENCE,
        data_target_average_len=[None] * len(buckets))

    dataset = data_io_pt.ParallelDataSet(*_get_random_bucketed_data(
        buckets, min_count=0, max_count=5, bucket_counts=bucket_counts))
    it = data_io_pt.ParallelSampleIter(dataset, buckets, batch_size,
                                       bucket_batch_sizes)

    with TemporaryDirectory() as work_dir:
        # Test 1
        it.next()
        expected_batch = it.next()

        fname = os.path.join(work_dir, "saved_iter")
        it.save_state(fname)

        it_loaded = data_io_pt.ParallelSampleIter(dataset, buckets, batch_size,
                                                  bucket_batch_sizes)
        it_loaded.reset()
        it_loaded.load_state(fname)
        loaded_batch = it_loaded.next()
        assert _data_batches_equal(expected_batch, loaded_batch)

        # Test 2
        it.reset()
        expected_batch = it.next()
        it.save_state(fname)

        it_loaded = data_io_pt.ParallelSampleIter(dataset, buckets, batch_size,
                                                  bucket_batch_sizes)
        it_loaded.reset()
        it_loaded.load_state(fname)

        loaded_batch = it_loaded.next()
        assert _data_batches_equal(expected_batch, loaded_batch)

        # Test 3
        it.reset()
        expected_batch = it.next()
        it.save_state(fname)
        it_loaded = data_io_pt.ParallelSampleIter(dataset, buckets, batch_size,
                                                  bucket_batch_sizes)
        it_loaded.reset()
        it_loaded.load_state(fname)

        loaded_batch = it_loaded.next()
        assert _data_batches_equal(expected_batch, loaded_batch)

        while it.iter_next():
            it.next()
            it_loaded.next()
        assert not it_loaded.iter_next()
Esempio n. 3
0
def test_get_batch_indices():
    max_bucket_size = 50
    batch_size = 10
    buckets = data_io_pt.define_parallel_buckets(100, 100, 10, True, 1.0)
    bucket_batch_sizes = data_io_pt.define_bucket_batch_sizes(
        buckets,
        batch_size,
        batch_type=C.BATCH_TYPE_SENTENCE,
        data_target_average_len=[None] * len(buckets))
    dataset = data_io_pt.ParallelDataSet(*_get_random_bucketed_data(
        buckets=buckets, min_count=1, max_count=max_bucket_size))

    indices = data_io_pt.get_batch_indices(
        dataset, bucket_batch_sizes=bucket_batch_sizes)

    # check for valid indices
    for buck_idx, start_pos in indices:
        assert 0 <= buck_idx < len(dataset)
        assert 0 <= start_pos < len(dataset.source[buck_idx]) - batch_size + 1

    # check that all indices are used for a filled-up dataset
    dataset = dataset.fill_up(bucket_batch_sizes)
    indices = data_io_pt.get_batch_indices(
        dataset, bucket_batch_sizes=bucket_batch_sizes)
    all_bucket_indices = set(list(range(len(dataset))))
    computed_bucket_indices = set([i for i, j in indices])

    assert not all_bucket_indices - computed_bucket_indices
Esempio n. 4
0
def test_parallel_data_set_permute():
    batch_size = 5
    buckets = data_io_pt.define_parallel_buckets(100, 100, 10, True, 1.0)
    bucket_batch_sizes = data_io_pt.define_bucket_batch_sizes(
        buckets,
        batch_size,
        batch_type=C.BATCH_TYPE_SENTENCE,
        data_target_average_len=[None] * len(buckets))
    dataset = data_io_pt.ParallelDataSet(*_get_random_bucketed_data(
        buckets, min_count=0, max_count=5)).fill_up(bucket_batch_sizes)

    permutations, inverse_permutations = data_io_pt.get_permutations(
        dataset.get_bucket_counts())

    assert len(permutations) == len(inverse_permutations) == len(dataset)
    dataset_restored = dataset.permute(permutations).permute(
        inverse_permutations)
    assert len(dataset) == len(dataset_restored)
    for buck_idx in range(len(dataset)):
        num_samples = dataset.source[buck_idx].shape[0]
        if num_samples:
            assert (dataset.source[buck_idx] ==
                    dataset_restored.source[buck_idx]).all()
            assert (dataset.target[buck_idx] ==
                    dataset_restored.target[buck_idx]).all()
        else:
            assert not dataset_restored.source[buck_idx].shape[0]
            assert not dataset_restored.target[buck_idx].shape[0]
Esempio n. 5
0
def test_sharded_and_parallel_iter_same_num_batches():
    """ Tests that a sharded data iterator with just a single shard produces as many shards as an iterator directly
    using the same dataset. """
    batch_size = 2
    num_batches_per_bucket = 10
    buckets = data_io_pt.define_parallel_buckets(100, 100, 10, True, 1.0)
    bucket_counts = [batch_size * num_batches_per_bucket for _ in buckets]
    num_batches = num_batches_per_bucket * len(buckets)
    bucket_batch_sizes = data_io_pt.define_bucket_batch_sizes(
        buckets,
        batch_size,
        batch_type=C.BATCH_TYPE_SENTENCE,
        data_target_average_len=[None] * len(buckets))

    dataset = data_io_pt.ParallelDataSet(*_get_random_bucketed_data(
        buckets, min_count=0, max_count=5, bucket_counts=bucket_counts))

    with TemporaryDirectory() as work_dir:
        shard_fname = os.path.join(work_dir, 'shard1')
        dataset.save(shard_fname)
        shard_fnames = [shard_fname]

        it_sharded = data_io_pt.ShardedParallelSampleIter(
            shard_fnames, buckets, batch_size, bucket_batch_sizes)

        it_parallel = data_io_pt.ParallelSampleIter(dataset, buckets,
                                                    batch_size,
                                                    bucket_batch_sizes)

        num_batches_seen = 0
        while it_parallel.iter_next():
            assert it_sharded.iter_next()
            it_parallel.next()
            it_sharded.next()
            num_batches_seen += 1
        assert num_batches_seen == num_batches

        print("Resetting...")
        it_sharded.reset()
        it_parallel.reset()

        num_batches_seen = 0
        while it_parallel.iter_next():
            assert it_sharded.iter_next()
            it_parallel.next()
            it_sharded.next()

            num_batches_seen += 1

        assert num_batches_seen == num_batches
Esempio n. 6
0
def test_parallel_data_set_fill_up():
    batch_size = 32
    buckets = data_io_pt.define_parallel_buckets(100, 100, 10, True, 1.0)
    bucket_batch_sizes = data_io_pt.define_bucket_batch_sizes(
        buckets,
        batch_size,
        batch_type=C.BATCH_TYPE_SENTENCE,
        data_target_average_len=[None] * len(buckets))
    dataset = data_io_pt.ParallelDataSet(
        *_get_random_bucketed_data(buckets, min_count=1, max_count=5))

    dataset_filled_up = dataset.fill_up(bucket_batch_sizes)
    assert len(dataset_filled_up.source) == len(dataset.source)
    assert len(dataset_filled_up.target) == len(dataset.target)
    for bidx in range(len(dataset)):
        bucket_batch_size = bucket_batch_sizes[bidx].batch_size
        assert dataset_filled_up.source[bidx].shape[0] == bucket_batch_size
        assert dataset_filled_up.target[bidx].shape[0] == bucket_batch_size
Esempio n. 7
0
def test_parallel_data_set():
    buckets = data_io_pt.define_parallel_buckets(100, 100, 10, True, 1.0)
    source, target = _get_random_bucketed_data(buckets,
                                               min_count=0,
                                               max_count=5)

    def check_equal(tensors1, tensors2):
        assert len(tensors1) == len(tensors2)
        for a1, a2 in zip(tensors1, tensors2):
            assert torch.equal(a1, a2)

    with TemporaryDirectory() as work_dir:
        dataset = data_io_pt.ParallelDataSet(source, target)
        fname = os.path.join(work_dir, 'dataset')
        dataset.save(fname)
        dataset_loaded = data_io_pt.ParallelDataSet.load(fname)
        check_equal(dataset.source, dataset_loaded.source)
        check_equal(dataset.target, dataset_loaded.target)