chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,497 @@
|
||||
import warnings
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import pyarrow
|
||||
import torch
|
||||
|
||||
from ray._common.utils import env_bool
|
||||
from ray.data.collate_fn import (
|
||||
TensorBatchReturnType,
|
||||
TensorBatchType,
|
||||
_is_nested_tensor_sequence,
|
||||
_is_tensor,
|
||||
_is_tensor_mapping,
|
||||
_is_tensor_sequence,
|
||||
_is_tensor_sequence_mapping,
|
||||
)
|
||||
from ray.data.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed
|
||||
|
||||
# Default non-blocking transfer for tensors.
|
||||
DEFAULT_TENSOR_NON_BLOCKING_TRANSFER = env_bool(
|
||||
"RAY_AIR_DEFAULT_TENSOR_NON_BLOCKING_TRANSFER",
|
||||
True,
|
||||
)
|
||||
|
||||
|
||||
def convert_table_to_torch_tensor(
|
||||
data_batch: pyarrow.Table,
|
||||
columns: Optional[Union[List[str], List[List[str]]]] = None,
|
||||
column_dtypes: Optional[Union[torch.dtype, List[torch.dtype]]] = None,
|
||||
unsqueeze: bool = True,
|
||||
) -> Union[torch.Tensor, List[torch.Tensor]]:
|
||||
"""Converts a PyArrow table to a torch Tensor or list of torch Tensors.
|
||||
|
||||
The return type matches the format of ``columns``: a flat list of column
|
||||
names produces a single tensor; a list of lists produces a list of tensors
|
||||
(one per group). If ``columns`` is None, all columns are used.
|
||||
|
||||
Args:
|
||||
data_batch: The PyArrow table to convert.
|
||||
columns: Column names to include. A list of lists returns one tensor
|
||||
per group (useful for multi-input models). None uses all columns.
|
||||
column_dtypes: Torch dtype(s) for the output. A single dtype applies
|
||||
to all columns/groups. A list must match the number of groups when
|
||||
``columns`` is a list of lists.
|
||||
unsqueeze: If True, reshape each per-column tensor from (N,) to (N, 1)
|
||||
before concatenating. Defaults to True.
|
||||
|
||||
Returns:
|
||||
A single tensor of shape (N, len(columns)), or a list of tensors when
|
||||
``columns`` is a list of lists.
|
||||
"""
|
||||
multi_input = columns and isinstance(columns[0], (list, tuple))
|
||||
|
||||
if columns is None:
|
||||
columns = data_batch.column_names
|
||||
|
||||
if not multi_input and column_dtypes and not isinstance(column_dtypes, torch.dtype):
|
||||
raise TypeError(
|
||||
"If `columns` is a list of strings, "
|
||||
"`column_dtypes` must be None or a single `torch.dtype`."
|
||||
f"Got {type(column_dtypes)} instead."
|
||||
)
|
||||
|
||||
if multi_input:
|
||||
if not isinstance(column_dtypes, (list, tuple)):
|
||||
column_dtypes = [column_dtypes] * len(columns)
|
||||
|
||||
return [
|
||||
_columns_to_tensor(data_batch, group, dtype, unsqueeze)
|
||||
for group, dtype in zip(columns, column_dtypes)
|
||||
]
|
||||
|
||||
return _columns_to_tensor(data_batch, columns, column_dtypes, unsqueeze)
|
||||
|
||||
|
||||
def _columns_to_tensor(
|
||||
table: pyarrow.Table,
|
||||
column_names: List[str],
|
||||
dtype: Optional[torch.dtype],
|
||||
unsqueeze: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Convert selected columns from a PyArrow table into a single tensor."""
|
||||
from ray.data._internal.arrow_block import ArrowBlockAccessor
|
||||
|
||||
numpy_batch = ArrowBlockAccessor(table).to_numpy(columns=column_names)
|
||||
|
||||
tensors = []
|
||||
for col in column_names:
|
||||
try:
|
||||
t = convert_ndarray_to_torch_tensor(numpy_batch[col], dtype)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Failed to convert column {col} to a Torch Tensor of dtype "
|
||||
f"{dtype}. See above exception chain for the exact failure."
|
||||
) from e
|
||||
|
||||
if unsqueeze:
|
||||
t = t.unsqueeze(1)
|
||||
|
||||
tensors.append(t)
|
||||
|
||||
if len(tensors) > 1:
|
||||
return torch.cat(tensors, dim=1)
|
||||
|
||||
return tensors[0]
|
||||
|
||||
|
||||
def convert_ndarray_to_torch_tensor(
|
||||
ndarray: np.ndarray,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
device: Optional[Union[str, "torch.device"]] = None,
|
||||
pin_memory: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Convert a NumPy ndarray to a Torch Tensor.
|
||||
|
||||
Args:
|
||||
ndarray: A NumPy ndarray that we wish to convert to a Torch Tensor.
|
||||
dtype: A Torch dtype for the created tensor; if None, the dtype will be
|
||||
inferred from the NumPy ndarray data.
|
||||
device: The device on which the tensor(s) should be placed; if None, the Torch
|
||||
tensor(s) will be constructed on the CPU.
|
||||
pin_memory: Whether to pin the memory of the created tensors.
|
||||
|
||||
Returns:
|
||||
A Torch Tensor.
|
||||
"""
|
||||
ndarray = _unwrap_ndarray_object_type_if_needed(ndarray)
|
||||
|
||||
# Object dtype cannot be converted into PyTorch Tensor.
|
||||
if ndarray.dtype.type is np.object_:
|
||||
raise RuntimeError(
|
||||
"Numpy array of object dtype cannot be converted to a Torch Tensor. This "
|
||||
"may because the numpy array is a ragged tensor--it contains items of "
|
||||
"different sizes. If using `iter_torch_batches()` API, you can pass in a "
|
||||
"`collate_fn` argument to specify custom logic to convert the Numpy array "
|
||||
"batch to a Torch tensor batch."
|
||||
)
|
||||
|
||||
# The numpy array is not always writeable as it can come from the Ray object store.
|
||||
# Numpy will throw a verbose warning here, which we suppress, as we don't write
|
||||
# to the tensors. We also don't want to copy the array to avoid memory overhead.
|
||||
# Original warning: https://github.com/pytorch/pytorch/blob/v1.13.0/
|
||||
# torch/csrc/utils/tensor_numpy.cpp#L198-L206
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore")
|
||||
result = torch.as_tensor(ndarray, dtype=dtype, device=device)
|
||||
|
||||
if pin_memory:
|
||||
assert result.device.type == "cpu", (
|
||||
"Pin memory is only supported for CPU tensors. "
|
||||
f"Got device: {result.device} and pin_memory: {pin_memory}."
|
||||
)
|
||||
result = result.pin_memory()
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def convert_ndarray_batch_to_torch_tensor_batch(
|
||||
ndarrays: Union[np.ndarray, Dict[str, np.ndarray]],
|
||||
dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None,
|
||||
device: Optional[Union[str, "torch.device"]] = None,
|
||||
pin_memory: bool = False,
|
||||
) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
|
||||
"""Convert a NumPy ndarray batch to a Torch Tensor batch.
|
||||
|
||||
Args:
|
||||
ndarrays: A (dict of) NumPy ndarray(s) that we wish to convert to a Torch Tensor.
|
||||
dtypes: A (dict of) Torch dtype(s) for the created tensor; if None, the dtype
|
||||
will be inferred from the NumPy ndarray data.
|
||||
device: The device on which the tensor(s) should be placed; if None, the Torch
|
||||
tensor(s) will be constructed on the CPU.
|
||||
pin_memory: Whether to pin the memory of the created tensors.
|
||||
|
||||
Returns:
|
||||
A (dict of) Torch Tensor(s).
|
||||
"""
|
||||
if isinstance(ndarrays, np.ndarray):
|
||||
# Single-tensor case.
|
||||
if isinstance(dtypes, dict):
|
||||
if len(dtypes) != 1:
|
||||
raise ValueError(
|
||||
"When constructing a single-tensor batch, only a single dtype "
|
||||
f"should be given, instead got: {dtypes}"
|
||||
)
|
||||
dtypes = next(iter(dtypes.values()))
|
||||
batch = convert_ndarray_to_torch_tensor(
|
||||
ndarrays,
|
||||
dtype=dtypes,
|
||||
device=device,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
else:
|
||||
# Multi-tensor case.
|
||||
batch = {
|
||||
col_name: convert_ndarray_to_torch_tensor(
|
||||
col_ndarray,
|
||||
dtype=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
|
||||
device=device,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
for col_name, col_ndarray in ndarrays.items()
|
||||
}
|
||||
|
||||
return batch
|
||||
|
||||
|
||||
def convert_ndarray_list_to_torch_tensor_list(
|
||||
ndarrays: Dict[str, List[np.ndarray]],
|
||||
dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None,
|
||||
device: Optional[Union[str, "torch.device"]] = None,
|
||||
pin_memory: bool = False,
|
||||
) -> Dict[str, List[torch.Tensor]]:
|
||||
"""Convert a dict mapping column names to lists of ndarrays to Torch Tensors.
|
||||
|
||||
Args:
|
||||
ndarrays: A dict mapping column names to lists of ndarrays that we wish to convert
|
||||
to Torch Tensors.
|
||||
dtypes: A (dict of) Torch dtype(s) for the created tensors; if None, the dtype
|
||||
will be inferred from the NumPy ndarray data.
|
||||
device: The device on which the tensor(s) should be placed; if None, the Torch
|
||||
tensor(s) will be constructed on the CPU.
|
||||
pin_memory: Whether to pin the memory of the created tensors.
|
||||
|
||||
Returns:
|
||||
A dict mapping column names to lists of Tensors.
|
||||
"""
|
||||
return {
|
||||
col_name: [
|
||||
convert_ndarray_batch_to_torch_tensor_batch(
|
||||
ndarray,
|
||||
dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
|
||||
device=device,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
for ndarray in col_ndarrays
|
||||
]
|
||||
for col_name, col_ndarrays in ndarrays.items()
|
||||
}
|
||||
|
||||
|
||||
def arrow_batch_to_tensors(
|
||||
batch: pyarrow.Table,
|
||||
dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None,
|
||||
combine_chunks: bool = False,
|
||||
pin_memory: bool = False,
|
||||
threadpool: Optional[ThreadPoolExecutor] = None,
|
||||
) -> Union[Dict[str, torch.Tensor], Dict[str, List[torch.Tensor]]]:
|
||||
"""Convert PyArrow batch to PyTorch tensors.
|
||||
|
||||
Args:
|
||||
batch: PyArrow batch to convert
|
||||
dtypes: A (dict of) Torch dtype(s) for the created tensors; if None, the dtype
|
||||
will be inferred from the NumPy ndarray data.
|
||||
combine_chunks: If True, combine chunks in Arrow batch before converting to
|
||||
tensors.
|
||||
pin_memory: Whether to pin the memory of the created tensors.
|
||||
threadpool: Optional ThreadPoolExecutor for parallel processing. If provided,
|
||||
columns/arrays will be processed in parallel. If None, processing is
|
||||
sequential.
|
||||
|
||||
Returns:
|
||||
When combine_chunks=True: A dictionary of column name to single tensor.
|
||||
When combine_chunks=False: A dictionary of column name to list of tensors.
|
||||
"""
|
||||
from ray.data._internal.arrow_block import ArrowBlockAccessor
|
||||
from ray.data._internal.arrow_ops import transform_pyarrow
|
||||
|
||||
if combine_chunks:
|
||||
numpy_batch = ArrowBlockAccessor(batch).to_batch_format("numpy")
|
||||
num_columns = len(numpy_batch)
|
||||
|
||||
if num_columns > 1 and threadpool is not None:
|
||||
# Process columns in parallel using provided threadpool
|
||||
def process_column(
|
||||
col_name_col_array: Tuple[str, np.ndarray]
|
||||
) -> Tuple[str, torch.Tensor]:
|
||||
col_name, col_array = col_name_col_array
|
||||
return col_name, convert_ndarray_batch_to_torch_tensor_batch(
|
||||
col_array,
|
||||
dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
|
||||
# Submit all columns to threadpool and collect results
|
||||
processed_cols = threadpool.map(process_column, numpy_batch.items())
|
||||
return dict(processed_cols)
|
||||
else:
|
||||
# Sequential processing for single column or single worker
|
||||
return {
|
||||
col_name: convert_ndarray_batch_to_torch_tensor_batch(
|
||||
col_array,
|
||||
dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
for col_name, col_array in numpy_batch.items()
|
||||
}
|
||||
else:
|
||||
numpy_list = transform_pyarrow.table_to_numpy_dict_chunked(
|
||||
batch,
|
||||
)
|
||||
# Count total number of arrays across all columns
|
||||
total_arrays = sum(len(arrays) for arrays in numpy_list.values())
|
||||
num_columns = len(numpy_list)
|
||||
|
||||
if total_arrays > 1 and threadpool is not None:
|
||||
# Process arrays in parallel using provided threadpool
|
||||
def process_array(
|
||||
array_item: Tuple[str, int, np.ndarray]
|
||||
) -> Tuple[str, int, torch.Tensor]:
|
||||
col_name, array_index, array = array_item
|
||||
return (
|
||||
col_name,
|
||||
array_index,
|
||||
convert_ndarray_batch_to_torch_tensor_batch(
|
||||
array,
|
||||
dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
|
||||
pin_memory=pin_memory,
|
||||
),
|
||||
)
|
||||
|
||||
# Flatten arrays with column name and index for parallel processing
|
||||
array_items = [
|
||||
(col_name, idx, array)
|
||||
for col_name, arrays in numpy_list.items()
|
||||
for idx, array in enumerate(arrays)
|
||||
]
|
||||
|
||||
# Submit all arrays to threadpool and collect results
|
||||
processed_arrays = list(threadpool.map(process_array, array_items))
|
||||
|
||||
# Initialize result with all columns from numpy_list, including empty ones
|
||||
# Pre-allocate lists of the correct size for each column
|
||||
result: Dict[str, List[torch.Tensor]] = {
|
||||
col_name: [None] * len(arrays)
|
||||
for col_name, arrays in numpy_list.items()
|
||||
}
|
||||
|
||||
# Populate result with processed tensors
|
||||
for col_name, array_index, tensor in processed_arrays:
|
||||
result[col_name][array_index] = tensor
|
||||
|
||||
return result
|
||||
else:
|
||||
# Sequential processing
|
||||
return convert_ndarray_list_to_torch_tensor_list(
|
||||
numpy_list,
|
||||
dtypes=dtypes,
|
||||
pin_memory=pin_memory,
|
||||
)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def concat_tensors_to_device(
|
||||
tensor_sequence: Sequence[torch.Tensor],
|
||||
device: Optional[Union[str, "torch.device"]] = None,
|
||||
non_blocking: bool = DEFAULT_TENSOR_NON_BLOCKING_TRANSFER,
|
||||
) -> torch.Tensor:
|
||||
"""Stack sequence of tensors into a contiguous GPU tensor.
|
||||
|
||||
Args:
|
||||
tensor_sequence: Sequence of tensors to stack
|
||||
device: The device to move tensors to. If None, tensors are not moved.
|
||||
non_blocking: If True, perform device transfer without forcing a
|
||||
synchronization.
|
||||
|
||||
Returns:
|
||||
A contiguous tensor on the target device
|
||||
"""
|
||||
# Assumes tensors have the same shape/dtype
|
||||
assert (
|
||||
tensor_sequence
|
||||
), f"Cannot stack empty sequence of tensors. Received: {tensor_sequence}"
|
||||
|
||||
assert all(
|
||||
isinstance(t, torch.Tensor) for t in tensor_sequence
|
||||
), "All items must be torch.Tensor. Found invalid types: " + str(
|
||||
[type(t) for t in tensor_sequence if not isinstance(t, torch.Tensor)]
|
||||
)
|
||||
|
||||
# If there is only one tensor and its device already matches, return it directly.
|
||||
if len(tensor_sequence) == 1 and (
|
||||
device is None or tensor_sequence[0].device == torch.device(device)
|
||||
):
|
||||
return tensor_sequence[0]
|
||||
|
||||
first_dtype = tensor_sequence[0].dtype
|
||||
assert all(t.dtype == first_dtype for t in tensor_sequence), (
|
||||
"All tensors must have the same dtype. "
|
||||
f"Expected: {first_dtype}, got: {[t.dtype for t in tensor_sequence]}"
|
||||
)
|
||||
|
||||
first_shape = tensor_sequence[0].shape[1:]
|
||||
assert all(t.shape[1:] == first_shape for t in tensor_sequence), (
|
||||
"All tensors must have the same shape[1:]. "
|
||||
f"Expected: {first_shape}, got: {[t.shape[1:] for t in tensor_sequence]}"
|
||||
)
|
||||
|
||||
first = tensor_sequence[0]
|
||||
dtype = first.dtype
|
||||
shape_tail = first.shape[1:]
|
||||
total_rows = sum(t.shape[0] for t in tensor_sequence)
|
||||
|
||||
# Allocate an empty Tensor on device
|
||||
result = torch.empty((total_rows, *shape_tail), dtype=dtype, device=device)
|
||||
|
||||
row_start = 0
|
||||
for t in tensor_sequence:
|
||||
row_end = row_start + t.shape[0]
|
||||
result[row_start:row_end].copy_(t, non_blocking=non_blocking)
|
||||
row_start = row_end
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _get_type_str(batch: Any) -> str:
|
||||
"""Get a string representation of the possibly nested type of the batch.
|
||||
|
||||
>>> import torch
|
||||
>>> _get_type_str([1, 2, "???"])
|
||||
'list[int | str]'
|
||||
>>> _get_type_str({"a": [1, 2, 3], "b": 4})
|
||||
'dict[str, int | list[int]]'
|
||||
>>> _get_type_str({"a": torch.tensor(1), "b": [torch.tensor(2)]})
|
||||
'dict[str, Tensor | list[Tensor]]'
|
||||
>>> _get_type_str({"a": torch.tensor(1), "b": {"c": torch.tensor(2)}})
|
||||
'dict[str, Tensor | dict[str, Tensor]]'
|
||||
"""
|
||||
curr_type = type(batch).__name__
|
||||
if isinstance(batch, (list, tuple)):
|
||||
val_types = " | ".join(sorted({_get_type_str(v) for v in batch}))
|
||||
invalid_type_str = f"{curr_type}[{val_types}]"
|
||||
elif isinstance(batch, dict):
|
||||
val_types = " | ".join(sorted({_get_type_str(v) for v in batch.values()}))
|
||||
invalid_type_str = f"{curr_type}[str, {val_types}]"
|
||||
else:
|
||||
invalid_type_str = curr_type
|
||||
return invalid_type_str
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def move_tensors_to_device(
|
||||
batch: TensorBatchType,
|
||||
device: Optional[Union[str, "torch.device"]] = None,
|
||||
non_blocking: bool = DEFAULT_TENSOR_NON_BLOCKING_TRANSFER,
|
||||
) -> TensorBatchReturnType:
|
||||
"""Move tensors to the specified device.
|
||||
|
||||
Concatenate nested lists/tuples of tensors along the first (batch) dimension.
|
||||
For example, for the input
|
||||
((feature_0_chunk_0,), (feature_1_chunk_0, feature_1_chunk_1))
|
||||
the output will be (feature_0_chunk_0, feature_1_chunk_0+1)
|
||||
where each feature is concatenated along the batch dimension.
|
||||
|
||||
Args:
|
||||
batch: A tensor or collection of tensors to move to device. Can be:
|
||||
- A single tensor
|
||||
- A sequence of tensors
|
||||
- A sequence of sequences of tensors. The inner sequence of tensors is
|
||||
combined during GPU transfer.
|
||||
- A mapping (e.g., dict) of keys to tensors or sequences of tensors. The
|
||||
sequence of tensors is combined during GPU transfer.
|
||||
device: The device to move tensors to. If None, tensors are not moved.
|
||||
non_blocking: If True, perform device transfer without forcing a
|
||||
synchronization.
|
||||
|
||||
Returns:
|
||||
The input tensors moved to the specified device
|
||||
"""
|
||||
if device is None:
|
||||
return batch
|
||||
|
||||
if _is_tensor(batch):
|
||||
return batch.to(device, non_blocking=non_blocking)
|
||||
elif _is_tensor_sequence(batch):
|
||||
return type(batch)([t.to(device, non_blocking=non_blocking) for t in batch])
|
||||
elif _is_nested_tensor_sequence(batch):
|
||||
return type(batch)(
|
||||
[concat_tensors_to_device(t, device, non_blocking) for t in batch]
|
||||
)
|
||||
elif _is_tensor_mapping(batch):
|
||||
return {k: t.to(device, non_blocking=non_blocking) for k, t in batch.items()}
|
||||
elif _is_tensor_sequence_mapping(batch):
|
||||
return {
|
||||
k: concat_tensors_to_device(v, device, non_blocking)
|
||||
for k, v in batch.items()
|
||||
}
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Invalid input type: {_get_type_str(batch)}.\n"
|
||||
"Expected one of the following: "
|
||||
"torch.Tensor, "
|
||||
"List/Tuple[torch.Tensor], "
|
||||
"Dict[str, torch.Tensor], "
|
||||
"Mapping[str, List/Tuple[torch.Tensor]]"
|
||||
)
|
||||
Reference in New Issue
Block a user