dmlc--dgl
Quan (Andy) Gan
a0e8cf0d3e
* fix uva sampling with features * fix * add is_listlike function to distinguish strings from sequence * fix Co-authored-by: nv-dlasalle <63612878+nv-dlasalle@users.noreply.github.com>
21 行
801 B
Python
21 行
801 B
Python
"""Feature storage for ``numpy.memmap`` object."""
|
|
import numpy as np
|
|
from .base import FeatureStorage, ThreadedFuture, register_storage_wrapper
|
|
from .. import backend as F
|
|
|
|
@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))
|