Classification of data's device and whether it is a batch. Based on data type determines if data should be treated as a batch and with which device. If the type can be recognized as a batch without being falsely categorized as such, it is. This includes lists of supported tensor-like ob
| 31 | |
| 32 | |
| 33 | class _Classification: |
| 34 | """Classification of data's device and whether it is a batch. |
| 35 | |
| 36 | Based on data type determines if data should be treated as a batch and with which device. |
| 37 | If the type can be recognized as a batch without being falsely categorized as such, it is. |
| 38 | This includes lists of supported tensor-like objects e.g. numpy arrays (the only list not |
| 39 | treated as a batch is a list of objects of primitive types), :class:`DataNodeDebug` and |
| 40 | TensorLists. |
| 41 | |
| 42 | Args: |
| 43 | data: Data to be classified. |
| 44 | type_name (str): Representation of argument type (input or keyword). |
| 45 | arg_constant_len (int): Only applicable for argument inputs that are of array type |
| 46 | (e.g. numpy array). If -1 does not modify the data. For positive value works like |
| 47 | `:func:types.Constant`, repeats the data `arg_constant_len` times. |
| 48 | """ |
| 49 | |
| 50 | def __init__(self, data, type_name, arg_constant_len=-1): |
| 51 | from nvidia.dali._debug_mode import DataNodeDebug |
| 52 | |
| 53 | is_batch, device, extracted = self._classify_data(data, type_name, arg_constant_len) |
| 54 | self.is_batch = is_batch |
| 55 | self.device = device |
| 56 | self.data = extracted |
| 57 | self.was_data_node = isinstance(data, DataNodeDebug) |
| 58 | self.original = data |
| 59 | |
| 60 | @staticmethod |
| 61 | def _classify_data(data, type_name, arg_constant_len): |
| 62 | """Returns tuple (is_batch, device, unpacked data).""" |
| 63 | from nvidia.dali._debug_mode import DataNodeDebug |
| 64 | |
| 65 | def is_primitive_type(x): |
| 66 | return isinstance(x, (int, float, bool, str)) |
| 67 | |
| 68 | def classify_array_input(arr): |
| 69 | if _types._is_numpy_array(arr): |
| 70 | device = "cpu" |
| 71 | elif _types._is_torch_tensor(arr): |
| 72 | device = "gpu" if arr.is_cuda else "cpu" |
| 73 | elif _types._is_mxnet_array(arr): |
| 74 | device = "gpu" if "gpu" in str(arr.context) else "cpu" |
| 75 | else: |
| 76 | raise RuntimeError(f"Unsupported array type '{type(arr)}'.") |
| 77 | |
| 78 | return False, device, arr |
| 79 | |
| 80 | def classify_array_kwarg(arr): |
| 81 | if _types._is_torch_tensor(arr): |
| 82 | if arr.is_cuda: |
| 83 | arr = arr.cpu().numpy() |
| 84 | elif _types._is_mxnet_array(arr): |
| 85 | import mxnet as mx |
| 86 | |
| 87 | if "gpu" in str(arr.context): |
| 88 | arr = arr.copyto(mx.cpu()) |
| 89 | elif not _types._is_numpy_array(arr): |
| 90 | raise RuntimeError(f"Unsupported array type '{type(arr)}'.") |
no outgoing calls
no test coverage detected