dmlc--dgl
Quan (Andy) Gan
067cd74420
* first commit * update * huh * fix * update * revert core * fix * update * rewrite * oops * address comments * add graph name * address comments * remove sample metadata file * address comments * fix * remove * add docs
24 行
847 B
Python
24 行
847 B
Python
import logging
|
|
import numpy as np
|
|
from numpy.lib.format import open_memmap
|
|
from .registry import register_array_parser
|
|
|
|
@register_array_parser("numpy")
|
|
class NumpyArrayParser(object):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def read(self, path):
|
|
logging.info('Reading from %s using numpy format' % path)
|
|
arr = np.load(path, mmap_mode='r')
|
|
logging.info('Done reading from %s' % path)
|
|
return arr
|
|
|
|
def write(self, path, arr):
|
|
logging.info('Writing to %s using numpy format' % path)
|
|
# np.save would load the entire memmap array up into CPU. So we manually open
|
|
# an empty npy file with memmap mode and manually flush it instead.
|
|
new_arr = open_memmap(path, mode='w+', dtype=arr.dtype, shape=arr.shape)
|
|
new_arr[:] = arr[:]
|
|
logging.info('Done writing to %s' % path)
|