| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869 |
- from collections import deque
- from typing import Any, List, Union
- import numpy as np
- from ray.rllib.utils.framework import try_import_tf, try_import_torch
- from ray.util.annotations import DeveloperAPI
- torch, _ = try_import_torch()
- _, tf, _ = try_import_tf()
- @DeveloperAPI
- def safe_isnan(value):
- """Check if a value is NaN.
- Args:
- value: The value to check.
- Returns:
- True if the value is NaN, False otherwise.
- """
- if torch and torch.is_tensor(value):
- return torch.isnan(value)
- if tf and tf.is_tensor(value):
- return tf.math.is_nan(value)
- return np.isnan(value)
- @DeveloperAPI
- def single_value_to_cpu(value):
- """Convert a single value to CPU if it's a tensor.
- TensorFlow tensors are always converted to numpy/python values.
- PyTorch tensors are converted to python scalars.
- """
- if torch and isinstance(value, torch.Tensor):
- return value.detach().cpu().item()
- elif tf and tf.is_tensor(value):
- return value.numpy()
- return value
- @DeveloperAPI
- def batch_values_to_cpu(values: Union[List[Any], deque]) -> List[Any]:
- """Convert a list or deque of GPU tensors to CPU scalars in a single operation.
- This function efficiently processes multiple PyTorch GPU tensors together by
- stacking them and performing a single .cpu() call. Assumes all values are either
- PyTorch tensors (on same device) or already CPU values.
- Args:
- values: A list or deque of values that may be GPU tensors.
- Returns:
- A list of CPU scalar values.
- """
- if not values:
- return []
- # Check if first value is a torch tensor - assume all are the same type
- if torch and isinstance(values[0], torch.Tensor):
- # Stack all tensors and move to CPU in one operation
- stacked = torch.stack(list(values))
- cpu_tensor = stacked.detach().cpu()
- return cpu_tensor.tolist()
- # Already CPU values
- return list(values)
|
