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