|
from typing import Any |
|
from ding.data.storage import Storage |
|
import pickle |
|
|
|
from ding.utils.file_helper import read_file, save_file |
|
|
|
|
|
class FileStorage(Storage): |
|
|
|
def save(self, data: Any) -> None: |
|
with open(self.path, "wb") as f: |
|
pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL) |
|
|
|
def load(self) -> Any: |
|
with open(self.path, "rb") as f: |
|
return pickle.load(f) |
|
|
|
|
|
class FileModelStorage(Storage): |
|
|
|
def save(self, state_dict: object) -> None: |
|
save_file(self.path, state_dict) |
|
|
|
def load(self) -> object: |
|
return read_file(self.path) |
|
|