wehub-resource-sync
70 lines
1.8 KiB
Python
70 lines
1.8 KiB
Python
"""For Tensor Serialization"""
|
|
from __future__ import absolute_import
|
|
|
|
from .. import backend as F
|
|
from .._ffi.function import _init_api
|
|
from ..ndarray import NDArray
|
|
|
|
__all__ = ["save_tensors", "load_tensors"]
|
|
|
|
_init_api("dgl.data.tensor_serialize")
|
|
|
|
|
|
def save_tensors(filename, tensor_dict):
|
|
"""
|
|
Save dict of tensors to file
|
|
|
|
Parameters
|
|
----------
|
|
filename : str
|
|
File name to store dict of tensors.
|
|
tensor_dict: dict of dgl NDArray or backend tensor
|
|
Python dict using string as key and tensor as value
|
|
|
|
Returns
|
|
----------
|
|
status : bool
|
|
Return whether save operation succeeds
|
|
"""
|
|
nd_dict = {}
|
|
is_empty_dict = len(tensor_dict) == 0
|
|
for key, value in tensor_dict.items():
|
|
if not isinstance(key, str):
|
|
raise Exception("Dict key has to be str")
|
|
if F.is_tensor(value):
|
|
nd_dict[key] = F.zerocopy_to_dgl_ndarray(value)
|
|
elif isinstance(value, NDArray):
|
|
nd_dict[key] = value
|
|
else:
|
|
raise Exception(
|
|
"Dict value has to be backend tensor or dgl ndarray"
|
|
)
|
|
|
|
return _CAPI_SaveNDArrayDict(filename, nd_dict, is_empty_dict)
|
|
|
|
|
|
def load_tensors(filename, return_dgl_ndarray=False):
|
|
"""
|
|
load dict of tensors from file
|
|
|
|
Parameters
|
|
----------
|
|
filename : str
|
|
File name to load dict of tensors.
|
|
return_dgl_ndarray: bool
|
|
Whether return dict of dgl NDArrays or backend tensors
|
|
|
|
Returns
|
|
---------
|
|
tensor_dict : dict
|
|
dict of tensor or ndarray based on return_dgl_ndarray flag
|
|
"""
|
|
nd_dict = _CAPI_LoadNDArrayDict(filename)
|
|
tensor_dict = {}
|
|
for key, value in nd_dict.items():
|
|
if return_dgl_ndarray:
|
|
tensor_dict[key] = value
|
|
else:
|
|
tensor_dict[key] = F.zerocopy_from_dgl_ndarray(value)
|
|
return tensor_dict
|