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)
(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',
)
| 59 | |
| 60 | |
| 61 | def _test(self, |
| 62 | datasets: T.List[TMPDataset], |
| 63 | dataset_ratios: T.List[float], |
| 64 | dataset_max_samples: T.List[int], |
| 65 | batch_size=10, |
| 66 | num_workers=0, |
| 67 | batch_wrt_length=True, |
| 68 | max_total_samples_per_epoch=-1, |
| 69 | max_batch_combined_size=-1, |
| 70 | batch_sampler_type='sort', |
| 71 | ): |
| 72 | """ |
| 73 | test plan: |
| 74 | - given created random datasets of different sizes |
| 75 | - create batch_preparer of different dataset ratios |
| 76 | - check the sum of the batches returned by the batch_preparer and the sum of the datasets (possibly weighted) |
| 77 | """ |
| 78 | |
| 79 | # create batch_preparer |
| 80 | batch_preparer = BatchPreparer( |
| 81 | batch_size=batch_size, |
| 82 | collate_fn=TMPCollate(), |
| 83 | num_workers=num_workers, |
| 84 | batch_wrt_length=batch_wrt_length, |
| 85 | shuffle=True, |
| 86 | max_total_samples_per_epoch=max_total_samples_per_epoch, |
| 87 | max_batch_combined_size=max_batch_combined_size, |
| 88 | batch_sampler_type=batch_sampler_type, |
| 89 | ) |
| 90 | # add datasets into batch_preparer |
| 91 | for i in range(len(datasets)): |
| 92 | batch_preparer.add_dataset(datasets[i], |
| 93 | ratio=dataset_ratios[i], |
| 94 | max_samples_per_epoch=dataset_max_samples[i]) |
| 95 | # test |
| 96 | # collect the batches |
| 97 | batches = [] |
| 98 | for i, batch in enumerate(batch_preparer): |
| 99 | batches.append(batch) |
| 100 | |
| 101 | # if all ratios are 1.0, we check if the sum is correct |
| 102 | if np.allclose(dataset_ratios, 1.0): |
| 103 | total_sum = np.sum([b.sum() for b in batches]) |
| 104 | total_sum_gt = 0 |
| 105 | for dset in datasets: |
| 106 | for i in range(len(dset)): |
| 107 | total_sum_gt += np.sum(dset[i]) |
| 108 | assert total_sum == total_sum_gt |
| 109 | else: |
| 110 | # make sure the dataset contains constant values |
| 111 | # check weight sum with ratio |
| 112 | total_sum = np.sum([b.sum() for b in batches]) |
| 113 | total_sum_gt = 0 |
| 114 | for didx in range(len(datasets)): |
| 115 | dset = datasets[didx] |
| 116 | if dataset_max_samples[didx] == -1: |
| 117 | dlen = len(dset) |
| 118 | else: |
no test coverage detected