dmlc--dgl
Hongzhi (Steve), Chen
3e5137feaf
Co-authored-by: Ubuntu <ubuntu@ip-172-31-28-63.ap-northeast-1.compute.internal>
26 行
826 B
Python
26 行
826 B
Python
"""Feature storage for ``numpy.memmap`` object."""
|
|
import numpy as np
|
|
|
|
from .. import backend as F
|
|
from .base import FeatureStorage, register_storage_wrapper, ThreadedFuture
|
|
|
|
|
|
@register_storage_wrapper(np.memmap)
|
|
class NumpyStorage(FeatureStorage):
|
|
"""FeatureStorage that asynchronously reads features from a ``numpy.memmap`` object."""
|
|
|
|
def __init__(self, arr):
|
|
self.arr = arr
|
|
|
|
# pylint: disable=unused-argument
|
|
def _fetch(self, indices, device, pin_memory=False):
|
|
result = F.zerocopy_from_numpy(self.arr[indices])
|
|
result = F.copy_to(result, device)
|
|
return result
|
|
|
|
# pylint: disable=unused-argument
|
|
def fetch(self, indices, device, pin_memory=False, **kwargs):
|
|
return ThreadedFuture(
|
|
target=self._fetch, args=(indices, device, pin_memory)
|
|
)
|