Copy contents of DALI tensor to PyTorch's Tensor. Parameters ---------- tensor_or_tl : TensorGPU or TensorListGPU arr : torch.Tensor Destination of the copy cuda_stream : torch.cuda.Stream, cudaStream_t or any value that can be cast to cudaStream_t.
(tensor_or_tl, arr, cuda_stream=None, non_blocking=False)
| 46 | |
| 47 | |
| 48 | def feed_ndarray(tensor_or_tl, arr, cuda_stream=None, non_blocking=False): |
| 49 | """ |
| 50 | Copy contents of DALI tensor to PyTorch's Tensor. |
| 51 | |
| 52 | Parameters |
| 53 | ---------- |
| 54 | tensor_or_tl : TensorGPU or TensorListGPU |
| 55 | arr : torch.Tensor |
| 56 | Destination of the copy |
| 57 | cuda_stream : torch.cuda.Stream, cudaStream_t or any value that can be cast to cudaStream_t. |
| 58 | CUDA stream to be used for the copy |
| 59 | (if not provided, an internal user stream will be selected) |
| 60 | In most cases, using pytorch's current stream is expected (for example, |
| 61 | if we are copying to a tensor allocated with torch.zeros(...)) |
| 62 | """ |
| 63 | dali_type = to_torch_type[tensor_or_tl.dtype] |
| 64 | if isinstance(tensor_or_tl, TensorListGPU): |
| 65 | dali_tensor = tensor_or_tl.as_tensor() |
| 66 | else: |
| 67 | dali_tensor = tensor_or_tl |
| 68 | |
| 69 | assert dali_type == arr.dtype, ( |
| 70 | f"The element type of DALI Tensor/TensorList " |
| 71 | f"doesn't match the element type of the target PyTorch Tensor: " |
| 72 | f"{dali_type} vs {arr.dtype}" |
| 73 | ) |
| 74 | |
| 75 | assert dali_tensor.shape() == list(arr.size()), ( |
| 76 | f"Shapes do not match: DALI tensor has size {dali_tensor.shape()}, " |
| 77 | f"but PyTorch Tensor has size {list(arr.size())}" |
| 78 | ) |
| 79 | cuda_stream = types._raw_cuda_stream_ptr(cuda_stream) |
| 80 | |
| 81 | # turn raw int to a c void pointer |
| 82 | c_type_pointer = ctypes.c_void_p(arr.data_ptr()) |
| 83 | tensor_or_tl.copy_to_external(c_type_pointer, cuda_stream, non_blocking) |
| 84 | return arr |
| 85 | |
| 86 | |
| 87 | def _test_copy_to_external(use_tensor_list, non_blocking): |
no test coverage detected