Load the model states and auxiliary states from disk. Usage: m = MyModel() m.compile(...) aux_states = m.load_states('mymodel.zip') Args: path: input file path (without the extension) Returns: dict
(self, fpath)
| 306 | os.rmdir(tmp_dir) |
| 307 | |
| 308 | def load_states(self, fpath): |
| 309 | """Load the model states and auxiliary states from disk. |
| 310 | |
| 311 | Usage: |
| 312 | m = MyModel() |
| 313 | m.compile(...) |
| 314 | aux_states = m.load_states('mymodel.zip') |
| 315 | |
| 316 | Args: |
| 317 | path: input file path (without the extension) |
| 318 | Returns: |
| 319 | dict |
| 320 | """ |
| 321 | |
| 322 | assert os.path.isfile(fpath), ( |
| 323 | "Failed to load states, %s is not exist." % fpath) |
| 324 | |
| 325 | timestamp = time.time() |
| 326 | tmp_dir = '/tmp/singa_load_states_%s' % timestamp |
| 327 | os.mkdir(tmp_dir) |
| 328 | |
| 329 | with zipfile.ZipFile(fpath, 'r') as zf: |
| 330 | zf.extractall(tmp_dir) |
| 331 | |
| 332 | tensor_dict_fp = tmp_dir + self.TENSOR_DICT_FILENAME |
| 333 | states_attr_fp = tmp_dir + self.STATES_ATTR_FILENAME |
| 334 | |
| 335 | with open(states_attr_fp) as f: |
| 336 | states_attr = json.load(f) |
| 337 | |
| 338 | tensor_dict = np.load(tensor_dict_fp) |
| 339 | |
| 340 | # restore singa tensor from numpy |
| 341 | model_states = dict() |
| 342 | aux_states = dict() |
| 343 | |
| 344 | for k in tensor_dict.files: |
| 345 | if states_attr[k]['state_type'] == self.MODEL_STATE_TYPE: |
| 346 | model_states[k] = tensor.from_numpy(tensor_dict[k]) |
| 347 | elif states_attr[k]['state_type'] == self.AUX_STATE_TYPE: |
| 348 | aux_states[k] = tensor.from_numpy(tensor_dict[k]) |
| 349 | |
| 350 | # restore model_states |
| 351 | self.set_states(model_states) |
| 352 | |
| 353 | # clean up tmp files |
| 354 | os.remove(tensor_dict_fp) |
| 355 | os.remove(states_attr_fp) |
| 356 | os.rmdir(tmp_dir) |
| 357 | return aux_states |
no test coverage detected