(data, batch_size, layout, device_id=None)
| 70 | |
| 71 | |
| 72 | def _prep_data_for_feed_input(data, batch_size, layout, device_id=None): |
| 73 | def to_numpy(x): |
| 74 | import numpy as np |
| 75 | |
| 76 | if _types._is_mxnet_array(x): |
| 77 | return x.asnumpy() |
| 78 | elif _types._is_torch_tensor(x): |
| 79 | return x.numpy() |
| 80 | else: |
| 81 | return np.asarray(x) |
| 82 | |
| 83 | # __cuda_array_interface__ doesn't provide any way to pass the information about the device |
| 84 | # where the memory is located. It is assumed that the current device is the one that |
| 85 | # the memory belongs to, unless the user sets the device explicitly |
| 86 | # creating TensorGPU/TensorListGPU |
| 87 | if isinstance(data, (_tensors.TensorListCPU, _tensors.TensorListGPU)): |
| 88 | if layout is not None: |
| 89 | _check_data_batch(data, batch_size, layout) |
| 90 | data = type(data)(data, layout) |
| 91 | elif isinstance(data, (list, tuple)): |
| 92 | inputs = [] |
| 93 | checked = False |
| 94 | for datum in data: |
| 95 | is_dlpack, is_gpu_data = _b.CheckDLPackCapsule(datum) |
| 96 | if not is_dlpack and not checked: |
| 97 | _check_data_batch(data, batch_size, layout) |
| 98 | checked = True |
| 99 | if isinstance(datum, (_tensors.TensorCPU, _tensors.TensorGPU)): |
| 100 | inp = type(datum)(datum, layout=layout) if layout is not None else datum |
| 101 | elif is_dlpack: |
| 102 | if is_gpu_data: |
| 103 | inp = _tensors.TensorGPU(datum, layout or "") |
| 104 | else: |
| 105 | inp = _tensors.TensorCPU(datum, layout or "") |
| 106 | elif hasattr(datum, "__cuda_array_interface__"): |
| 107 | array_device_id = _types._get_device_id_for_array(datum) |
| 108 | if array_device_id is None: |
| 109 | array_device_id = device_id |
| 110 | inp = _tensors.TensorGPU(datum, layout or "", array_device_id) |
| 111 | else: |
| 112 | datum = to_numpy(datum) |
| 113 | inp = _tensors.TensorCPU(datum, layout or "") |
| 114 | inputs.append(inp) |
| 115 | assert all( |
| 116 | isinstance(inp, type(inputs[0])) for inp in inputs |
| 117 | ), "Mixed input types are not support, all need to reside on the CPU or GPU" |
| 118 | data = inputs |
| 119 | else: |
| 120 | is_dlpack, is_gpu_data = _b.CheckDLPackCapsule(data) |
| 121 | if not is_dlpack: |
| 122 | _check_data_batch(data, batch_size, layout) |
| 123 | if hasattr(data, "__cuda_array_interface__"): |
| 124 | array_device_id = _types._get_device_id_for_array(data) |
| 125 | if array_device_id is None: |
| 126 | array_device_id = device_id |
| 127 | data = _tensors.TensorListGPU(data, layout or "", array_device_id) |
| 128 | elif is_dlpack: |
| 129 | if is_gpu_data: |
no test coverage detected