ding.data.storage.file

ding.data.storage.file

Full Source Code

../ding/data/storage/file.py

1from typing import Any 2from ding.data.storage import Storage 3import pickle 4 5from ding.utils.file_helper import read_file, save_file 6 7 8class FileStorage(Storage): 9 10 def save(self, data: Any) -> None: 11 with open(self.path, "wb") as f: 12 pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL) 13 14 def load(self) -> Any: 15 with open(self.path, "rb") as f: 16 return pickle.load(f) 17 18 19class FileModelStorage(Storage): 20 21 def save(self, state_dict: object) -> None: 22 save_file(self.path, state_dict) 23 24 def load(self) -> object: 25 return read_file(self.path)