wehub-resource-sync
59 lines
2.1 KiB
Python
59 lines
2.1 KiB
Python
"""Feature storages for PyTorch tensors."""
|
|
|
|
import torch
|
|
|
|
from ..utils import gather_pinned_tensor_rows
|
|
from .base import register_storage_wrapper
|
|
from .tensor import BaseTensorStorage
|
|
|
|
|
|
def _fetch_cpu(indices, tensor, feature_shape, device, pin_memory, **kwargs):
|
|
result = torch.empty(
|
|
indices.shape[0],
|
|
*feature_shape,
|
|
dtype=tensor.dtype,
|
|
pin_memory=pin_memory,
|
|
)
|
|
torch.index_select(tensor, 0, indices, out=result)
|
|
kwargs["non_blocking"] = pin_memory
|
|
result = result.to(device, **kwargs)
|
|
return result
|
|
|
|
|
|
def _fetch_cuda(indices, tensor, device, **kwargs):
|
|
return torch.index_select(tensor, 0, indices).to(device, **kwargs)
|
|
|
|
|
|
@register_storage_wrapper(torch.Tensor)
|
|
class PyTorchTensorStorage(BaseTensorStorage):
|
|
"""Feature storages for slicing a PyTorch tensor."""
|
|
|
|
def fetch(self, indices, device, pin_memory=False, **kwargs):
|
|
device = torch.device(device)
|
|
storage_device_type = self.storage.device.type
|
|
indices_device_type = indices.device.type
|
|
if storage_device_type != "cuda":
|
|
if indices_device_type == "cuda":
|
|
if self.storage.is_pinned():
|
|
return gather_pinned_tensor_rows(self.storage, indices)
|
|
else:
|
|
raise ValueError(
|
|
f"Got indices on device {indices.device} whereas the feature tensor "
|
|
f"is on {self.storage.device}. Please either (1) move the graph "
|
|
f"to GPU with to() method, or (2) pin the graph with "
|
|
f"pin_memory_() method."
|
|
)
|
|
# CPU to CPU or CUDA - use pin_memory and async transfer if possible
|
|
else:
|
|
return _fetch_cpu(
|
|
indices,
|
|
self.storage,
|
|
self.storage.shape[1:],
|
|
device,
|
|
pin_memory,
|
|
**kwargs,
|
|
)
|
|
else:
|
|
# CUDA to CUDA or CPU
|
|
return _fetch_cuda(indices, self.storage, device, **kwargs)
|