(
reduce_fn, batch_fn, keep_dims, axes, output_type, add_mean_input=False, ddof=0, layout=None
)
| 131 | |
| 132 | |
| 133 | def run_dali( |
| 134 | reduce_fn, batch_fn, keep_dims, axes, output_type, add_mean_input=False, ddof=0, layout=None |
| 135 | ): |
| 136 | batch_size = batch_fn.batch_size() |
| 137 | |
| 138 | # Needed due to how ExternalSource API works. It fails on methods, partials. |
| 139 | def get_batch(): |
| 140 | return batch_fn() |
| 141 | |
| 142 | result_cpu = [] |
| 143 | result_gpu = [] |
| 144 | |
| 145 | pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0) |
| 146 | |
| 147 | args = {"keep_dims": keep_dims, "axes": axes} |
| 148 | if output_type is not None: |
| 149 | args["dtype"] = np_type_to_dali(output_type) |
| 150 | |
| 151 | with pipe: |
| 152 | input = fn.external_source(source=get_batch, layout=layout) |
| 153 | if not add_mean_input: |
| 154 | reduced_cpu = reduce_fn(input, **args) |
| 155 | reduced_gpu = reduce_fn(input.gpu(), **args) |
| 156 | else: |
| 157 | mean = fn.reductions.mean(input, **args) |
| 158 | args["ddof"] = ddof |
| 159 | reduced_cpu = reduce_fn(input, mean, **args) |
| 160 | reduced_gpu = reduce_fn(input.gpu(), mean.gpu(), **args) |
| 161 | pipe.set_outputs(reduced_cpu, reduced_gpu) |
| 162 | |
| 163 | for _ in range(batch_fn.num_iter()): |
| 164 | output = pipe.run() |
| 165 | check_layout(output[0], layout, axes, keep_dims) |
| 166 | check_layout(output[1], layout, axes, keep_dims) |
| 167 | reduced_cpu = output[0].as_array() |
| 168 | reduced_gpu = output[1].as_cpu().as_array() |
| 169 | result_cpu.append(reduced_cpu) |
| 170 | result_gpu.append(reduced_gpu) |
| 171 | |
| 172 | return result_cpu, result_gpu |
| 173 | |
| 174 | |
| 175 | def run_numpy(reduce_fn, batch_fn, keep_dims, axes, output_type, ddof=None): |
no test coverage detected