def to_numpy(x):
if types._is_mxnet_array(x):
return x.asnumpy()
elif types._is_torch_tensor(x):
return x.numpy()
else:
return x
def _get_default_stream_for_array(array):
if types._is_torch_tensor(array):
import torch
return torch.cuda.current_stream().cuda_stream
elif types._is_cupy_array(array):
import cupy
return cupy.cuda.get_current_stream().ptr
else:
return None
def classify_array_input(arr):
if _types._is_numpy_array(arr):
device = 'cpu'
elif _types._is_torch_tensor(arr):
device = 'gpu' if arr.is_cuda else 'cpu'
elif _types._is_mxnet_array(arr):
device = 'gpu' if 'gpu' in str(arr.context) else 'cpu'
else:
raise RuntimeError(f"Unsupported array type '{type(arr)}'.")
return False, device, arr
def classify_array_kwarg(arr):
if _types._is_torch_tensor(arr):
if arr.is_cuda:
arr = arr.cpu().numpy()
elif _types._is_mxnet_array(arr):
import mxnet as mx
if 'gpu' in str(arr.context):
arr = arr.copyto(mx.cpu())
elif not _types._is_numpy_array(arr):
raise RuntimeError(f"Unsupported array type '{type(arr)}'.")
arr = _types._preprocess_constant_array_type(arr)
arr = _tensors.TensorListCPU([_tensors.TensorCPU(arr)] * arg_constant_len)
return True, 'cpu', arr
def _to_numpy(sample):
if isinstance(sample, np.ndarray):
return sample
if types._is_mxnet_array(sample):
if sample.ctx.device_type != 'cpu':
raise TypeError(
"GPU tensors are not supported/ Got an MXNet GPU tensor.")
return sample.asnumpy()
if types._is_torch_tensor(sample):
if sample.device.type != 'cpu':
raise TypeError(
"GPU tensors are not supported. Got a PyTorch GPU tensor.")
return sample.numpy()
elif isinstance(sample, tensors.TensorCPU):
return np.array(sample)
raise TypeError(
"Unsupported callback return type. Expected NumPy array, PyTorch or MXNet cpu tensors, "
"DALI TensorCPU, or list or tuple of them.")
def assert_cpu_sample_data_type(
sample, error_str="Unsupported callback return type. Got: `{}`."):
import_numpy()
if isinstance(sample, np.ndarray):
return True
if types._is_mxnet_array(sample):
if sample.ctx.device_type != 'cpu':
raise TypeError(
"Unsupported callback return type. "
"GPU tensors are not supported. Got an MXNet GPU tensor.")
return True
if types._is_torch_tensor(sample):
if sample.device.type != 'cpu':
raise TypeError(
"Unsupported callback return type. "
"GPU tensors are not supported. Got a PyTorch GPU tensor.")
return True
elif isinstance(sample, tensors.TensorCPU):
return True
raise TypeError(error_str.format(type(sample)))
def _classify_data(data):
""" Returns tuple (use_external_source, device, unpacked data).
Based on data type determines if we should use external_source 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 objects e.g. numpy arrays (the only list not treated as a batch is a list of objects of
primitive types), :class:`DataNodeDebug` and TensorLists.
"""
def is_primitive_type(x):
return isinstance(x, (int, float, bool, str))
if isinstance(data, list):
if any([is_primitive_type(d) for d in data]):
return False, 'cpu', data
device = None
data_list = []
for d in data:
_, cur_device, val = _PipelineDebug._classify_data(d)
if device is None:
device = cur_device
data_list.append(val)
return True, device, data_list
else:
if is_primitive_type(data) or _types._is_numpy_array(data) or \
_types._is_mxnet_array(data) or isinstance(data, _tensors.TensorCPU):
return False, 'cpu', data
if _types._is_torch_tensor(data):
return False, 'gpu' if data.is_cuda else 'cpu', data
if hasattr(data, '__cuda_array_interface__') or isinstance(
data, _tensors.TensorGPU):
return False, 'gpu', data
if isinstance(data, DataNodeDebug):
return True, data.device, data.get()
if isinstance(data, _tensors.TensorListCPU):
return True, 'cpu', data
if isinstance(data, _tensors.TensorListGPU):
return True, 'gpu', data
return False, 'cpu', data
def _assert_valid_data_type(sample):
import_numpy()
if isinstance(sample, np.ndarray):
return
if types._is_mxnet_array(sample):
if sample.ctx.device_type != 'cpu':
raise TypeError(
"Unsupported callback return type. "
"GPU tensors are not supported. Got an MXNet GPU tensor.")
return
if types._is_torch_tensor(sample):
if sample.device.type != 'cpu':
raise TypeError(
"Unsupported callback return type. "
"GPU tensors are not supported. Got a PyTorch GPU tensor.")
return
elif isinstance(sample, tensors.TensorCPU):
return
raise TypeError((
"Unsupported callback return type. Expected NumPy array, PyTorch or MXNet cpu tensors, "
"DALI TensorCPU, or list or tuple of them. Got `{}` instead.").format(
type(sample)))
