Split a batch tensor into a list of individual tensors using zero-copy. Args: batch_tensor: The batch tensor to split. Returns: A list of individual tensors.
(batch_tensor: cvcuda.Tensor)
| 146 | |
| 147 | |
| 148 | def zero_copy_split(batch_tensor: cvcuda.Tensor) -> list[cvcuda.Tensor]: |
| 149 | """ |
| 150 | Split a batch tensor into a list of individual tensors using zero-copy. |
| 151 | |
| 152 | Args: |
| 153 | batch_tensor: The batch tensor to split. |
| 154 | |
| 155 | Returns: |
| 156 | A list of individual tensors. |
| 157 | """ |
| 158 | # Helper object which has a __cuda_array_interface__ |
| 159 | class CudaBuffer: |
| 160 | __cuda_array_interface__ = None |
| 161 | obj = None |
| 162 | |
| 163 | # The high-level overview of this function is to create a custom 'view' |
| 164 | # (using PyTorch terminology) for each individual Tensor in the batch. |
| 165 | # We will do this by creating a custom CUDA buffer with a __cuda_array_interface__ |
| 166 | # setup from the batch tensor's CUDA buffer. |
| 167 | # The steps are: |
| 168 | # 1. Get the CUDA buffer from the batch tensor |
| 169 | # 2. For each image in the batch, create a custom CUDA buffer with a __cuda_array_interface__ |
| 170 | # setup from the batch tensor's CUDA buffer. |
| 171 | # 3. Create a new Tensor from the custom CUDA buffer. |
| 172 | |
| 173 | cuda_interface = batch_tensor.cuda().__cuda_array_interface__ |
| 174 | batch_dim = batch_tensor.shape[0] |
| 175 | height, width, channels = batch_tensor.shape[1:] |
| 176 | dtype = batch_tensor.dtype |
| 177 | |
| 178 | # Get strides from the tensor to handle padded memory correctly |
| 179 | strides = cuda_interface.get("strides") |
| 180 | if strides is None: |
| 181 | # If no strides, assume C-contiguous layout |
| 182 | batch_stride_bytes = height * width * channels * dtype.itemsize |
| 183 | item_strides = None |
| 184 | else: |
| 185 | # Use the actual batch stride from the tensor |
| 186 | batch_stride_bytes = strides[0] |
| 187 | # Preserve the HWC strides for the resulting tensors |
| 188 | item_strides = strides[1:] |
| 189 | |
| 190 | new_tensors: list[cvcuda.Tensor] = [] |
| 191 | for i in range(batch_dim): |
| 192 | offset_ptr = cuda_interface["data"][0] + i * batch_stride_bytes |
| 193 | |
| 194 | offset_buffer = CudaBuffer() |
| 195 | buffer_interface = { |
| 196 | "shape": (height, width, channels), |
| 197 | "typestr": dtype.str, |
| 198 | "data": (offset_ptr, False), |
| 199 | "version": 3, |
| 200 | } |
| 201 | |
| 202 | # Include strides if the tensor has non-contiguous memory layout |
| 203 | if item_strides is not None: |
| 204 | buffer_interface["strides"] = item_strides |
| 205 |
no test coverage detected