dmlc--dgl
peizhou001
27 行
850 B
Python
27 行
850 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)
|