def _SaveVarDict( path: str, var_dict: Optional[Dict[str, Union[FileBackendVariableBlob, EagerBlobTrait]]] = None, ) -> None: if var_dict is None: var_dict = GetAllVariables() def IsFileOrNonEmptyDir(path): if os.path.isfile(path): return True if os.path.isdir(path) and len(os.listdir(path)) != 0: return True return False assert not IsFileOrNonEmptyDir( path ), "{} is a file or non-empty directory! Note that flow.save is different from torch.save. It saves each weight as a separated file so that a directory instead of a file should be given.".format( path) os.makedirs(path, exist_ok=True) for (name, var) in var_dict.items(): meta_info = variable_meta_info_pb.VariableMetaInfo() meta_info.shape.dim[:] = var.shape meta_info.data_type = oneflow._oneflow_internal.deprecated.GetProtoDtype4OfDtype( var.dtype) var_dir = os.path.join(path, name) param_path = os.path.join(var_dir, DATA_FILENAME) os.makedirs(os.path.dirname(param_path)) with open(param_path, "wb") as f: for (_, _, slice) in _ReadSlice(var): f.write(slice.tobytes()) with open(os.path.join(var_dir, META_INFO_FILENAME), "w") as f: f.write(text_format.MessageToString(meta_info)) with open(os.path.join(path, "snapshot_done"), "w"): pass
def _save_tensor_to_disk(tensor: "oneflow.Tensor", dir_name: Union[str, Path]) -> None: os.makedirs(dir_name, exist_ok=True) meta_info = variable_meta_info_pb.VariableMetaInfo() meta_info.shape.dim[:] = tensor.shape meta_info.data_type = oneflow._oneflow_internal.deprecated.GetProtoDtype4OfDtype( tensor.dtype) data_path = os.path.join(dir_name, DATA_FILENAME) with open(data_path, "wb") as f: f.write(tensor.numpy().tobytes()) with open(os.path.join(dir_name, META_INFO_FILENAME), "w") as f: f.write(text_format.MessageToString(meta_info))
def __init__( self, var_dir: str, dtype: Optional[dtype_util.dtype] = None, shape: Optional[Sequence[int]] = None, ): data_path = os.path.join(var_dir, DATA_FILENAME) assert os.path.isfile(data_path) self.var_dir_ = var_dir meta_info_path = os.path.join(self.var_dir_, META_INFO_FILENAME) if os.path.exists(meta_info_path): meta_info = variable_meta_info_pb.VariableMetaInfo() with open(meta_info_path) as f: text_format.Parse(f.read(), meta_info) self.has_meta_info_ = True else: self.has_meta_info_ = False if self.has_meta_info_: assert dtype is None and shape is None self.shape_ = tuple(meta_info.shape.dim) self.dtype_ = dtype_util.convert_proto_dtype_to_oneflow_dtype( meta_info.data_type) else: if shape is not None and dtype is not None: self.shape_ = shape self.dtype_ = dtype self.has_meta_info_ = True elif shape is not None or dtype is not None: raise RuntimeError( "both or neither of shape and dtype should be None") else: pass if self.has_meta_info_: itemsize = np.dtype( dtype_util.convert_oneflow_dtype_to_numpy_dtype( self.dtype_)).itemsize assert os.path.getsize(data_path) == np.prod( self.shape).item() * itemsize
def SaveVarDict( path: str, var_dict: Optional[Dict[str, Union[FileBackendVariableBlob, EagerBlobTrait]]] = None, ) -> None: """ Save `var_dict` to `path` """ oneflow.sync_default_session() if var_dict is None: var_dict = GetAllVariables() def IsFileOrNonEmptyDir(path): if os.path.isfile(path): return True if os.path.isdir(path) and len(os.listdir(path)) != 0: return True return False assert not IsFileOrNonEmptyDir( path), "Non-empty directory {} already exists!".format(path) os.makedirs(path, exist_ok=True) for name, var in var_dict.items(): meta_info = variable_meta_info_pb.VariableMetaInfo() meta_info.shape.dim[:] = var.shape meta_info.data_type = oneflow_api.deprecated.GetProtoDtype4OfDtype( var.dtype) var_dir = os.path.join(path, name) param_path = os.path.join(var_dir, DATA_FILENAME) os.makedirs(os.path.dirname(param_path)) with open(param_path, "wb") as f: for _, _, slice in _ReadSlice(var): f.write(slice.tobytes()) with open(os.path.join(var_dir, META_INFO_FILENAME), "w") as f: f.write(text_format.MessageToString(meta_info)) # write a empty file 'snapshot_done', indicating that # the save process finishes normally with open(os.path.join(path, "snapshot_done"), "w"): pass