| 20 | |
| 21 | |
| 22 | def test_select_impls(): |
| 23 | # Test recursive select that returns tuple of outputs (due to the ops having 2 outputs). |
| 24 | # Without supporting nested structures it encountered a ((DataNode, DataNode),) branch output |
| 25 | # and crashed. |
| 26 | |
| 27 | def _select_fwd(op_range_lo, op_range_hi, ops, selected_op_idx, op_args, op_kwargs): |
| 28 | assert op_range_lo <= op_range_hi |
| 29 | if op_range_lo == op_range_hi: |
| 30 | return ops[op_range_lo](*op_args, **op_kwargs) |
| 31 | mid = (op_range_lo + op_range_hi) // 2 |
| 32 | if selected_op_idx <= mid: |
| 33 | ret = _select_fwd(op_range_lo, mid, ops, selected_op_idx, op_args, op_kwargs) |
| 34 | else: |
| 35 | ret = _select_fwd(mid + 1, op_range_hi, ops, selected_op_idx, op_args, op_kwargs) |
| 36 | return ret |
| 37 | |
| 38 | def _select_unpack(op_range_lo, op_range_hi, ops, selected_op_idx, op_args, op_kwargs): |
| 39 | assert op_range_lo <= op_range_hi |
| 40 | if op_range_lo == op_range_hi: |
| 41 | return ops[op_range_lo](*op_args, **op_kwargs) |
| 42 | mid = (op_range_lo + op_range_hi) // 2 |
| 43 | if selected_op_idx <= mid: |
| 44 | a, b = _select_unpack(op_range_lo, mid, ops, selected_op_idx, op_args, op_kwargs) |
| 45 | else: |
| 46 | a, b = _select_unpack(mid + 1, op_range_hi, ops, selected_op_idx, op_args, op_kwargs) |
| 47 | return a, b |
| 48 | |
| 49 | def select(ops, selected_op_idx, *op_args, unpacking_select=False, **op_kwargs): |
| 50 | if unpacking_select: |
| 51 | return _select_unpack(0, len(ops) - 1, ops, selected_op_idx, op_args, op_kwargs) |
| 52 | else: |
| 53 | return _select_fwd(0, len(ops) - 1, ops, selected_op_idx, op_args, op_kwargs) |
| 54 | |
| 55 | def rotate(image, label): |
| 56 | image = fn.rotate(image, angle=42) |
| 57 | return image, label |
| 58 | |
| 59 | def color(image, label): |
| 60 | image = fn.color_twist(image, saturation=0) |
| 61 | return image, label |
| 62 | |
| 63 | @pipeline_def(enable_conditionals=True, num_threads=4, batch_size=8, device_id=0) |
| 64 | def pipeline(unpacking_select): |
| 65 | image = types.Constant(np.full((200, 300, 3), 42, dtype=np.uint8), device="cpu") |
| 66 | label = types.Constant(np.array(1), device="cpu") |
| 67 | ops = [rotate, color] |
| 68 | op_idx = fn.random.uniform(values=list(range(len(ops)))) |
| 69 | image, label = select( |
| 70 | ops, op_idx, image=image, label=label, unpacking_select=unpacking_select |
| 71 | ) |
| 72 | return image, label |
| 73 | |
| 74 | pipe_unpacking = pipeline(unpacking_select=True) |
| 75 | pipe_unpacking.run() |
| 76 | |
| 77 | pipe_forwarding = pipeline(unpacking_select=False) |
| 78 | pipe_forwarding.run() |
| 79 | |