Class MyTestCase

tests/cdslib/data/dataloader/test_batch_preparer.py:45–224

class MyTestCase(unittest.TestCase):

    @staticmethod
    def _create_dataset(total_samples, seq_len, dim, val=None):
        """
        Return a pytorch dataset instance, which contains (total_samples, dim).
        :param val: constant fill in value. None: random array
        """
        if val is None:
            arr = np.random.randn(total_samples, seq_len, dim)
        else:
            arr = np.ones((total_samples, seq_len, dim)) * val
        arr = arr.astype(dtype='int32')
        return TMPDataset(arr)


    def _test(self,
              datasets: T.List[TMPDataset],
              dataset_ratios: T.List[float],
              dataset_max_samples: T.List[int],
              batch_size=10,
              num_workers=0,
              batch_wrt_length=True,
              max_total_samples_per_epoch=-1,
              max_batch_combined_size=-1,
              batch_sampler_type='sort',
              ):
        """
        test plan:
        - given created random datasets of different sizes
        - create batch_preparer of different dataset ratios
        - check the sum of the batches returned by the batch_preparer and the sum of the datasets (possibly weighted)
        """

        # create batch_preparer
        batch_preparer = BatchPreparer(
            batch_size=batch_size,
            collate_fn=TMPCollate(),
            num_workers=num_workers,
            batch_wrt_length=batch_wrt_length,
            shuffle=True,
            max_total_samples_per_epoch=max_total_samples_per_epoch,
            max_batch_combined_size=max_batch_combined_size,
            batch_sampler_type=batch_sampler_type,
        )
        # add datasets into batch_preparer
        for i in range(len(datasets)):
            batch_preparer.add_dataset(datasets[i],
                                       ratio=dataset_ratios[i],
                                       max_samples_per_epoch=dataset_max_samples[i])
        # test
        # collect the batches
        batches = []
        for i, batch in enumerate(batch_preparer):
            batches.append(batch)

        # if all ratios are 1.0, we check if the sum is correct
        if np.allclose(dataset_ratios, 1.0):