| 120 | |
| 121 | |
| 122 | def _test_memory_consumption(device, test_case): |
| 123 | batch_size = 32 |
| 124 | num_iters = 128 |
| 125 | |
| 126 | if device == "cpu": |
| 127 | import numpy as np |
| 128 | |
| 129 | fw = np |
| 130 | else: |
| 131 | fw = cp |
| 132 | |
| 133 | def no_copy_sample(): |
| 134 | batch = [fw.full((1024, 1024, 4), i, dtype=fw.int32) for i in range(batch_size)] |
| 135 | |
| 136 | def cb(sample_info): |
| 137 | return batch[sample_info.idx_in_batch] |
| 138 | |
| 139 | return cb |
| 140 | |
| 141 | def copy_sample(): |
| 142 | def cb(sample_info): |
| 143 | return fw.full((1024, 1024, 4), sample_info.idx_in_batch, dtype=fw.int32) |
| 144 | |
| 145 | return cb |
| 146 | |
| 147 | def copy_batch(): |
| 148 | def cb(): |
| 149 | return fw.full((batch_size, 1024, 1024, 4), 42, dtype=fw.int32) |
| 150 | |
| 151 | return cb |
| 152 | |
| 153 | cases = { |
| 154 | "no_copy_sample": (no_copy_sample, True, False), |
| 155 | "copy_sample": (copy_sample, False, False), |
| 156 | "copy_batch": (copy_batch, False, True), |
| 157 | } |
| 158 | |
| 159 | cb, no_copy, batch_mode = cases[test_case] |
| 160 | |
| 161 | @pipeline_def |
| 162 | def pipeline(): |
| 163 | return fn.external_source(source=cb(), device=device, batch=batch_mode, no_copy=no_copy) |
| 164 | |
| 165 | pipe = pipeline(batch_size=batch_size, num_threads=4, device_id=0) |
| 166 | for _ in range(num_iters): |
| 167 | pipe.run() |
| 168 | |
| 169 | |
| 170 | def test_memory_consumption(): |