Save states. Args: fpath: output file path (without the extension) aux_states(dict): values are standard data types or Tensor, e.g., epoch ID, learning rate, optimizer states
(self, fpath, aux_states={})
| 242 | return self.forward(*input, **kwargs) |
| 243 | |
| 244 | def save_states(self, fpath, aux_states={}): |
| 245 | """Save states. |
| 246 | |
| 247 | Args: |
| 248 | fpath: output file path (without the extension) |
| 249 | aux_states(dict): values are standard data types or Tensor, |
| 250 | e.g., epoch ID, learning rate, optimizer states |
| 251 | """ |
| 252 | assert not os.path.isfile(fpath), ( |
| 253 | "Failed to save states, %s is already existed." % fpath) |
| 254 | |
| 255 | states = self.get_states() |
| 256 | |
| 257 | # save states data and attr |
| 258 | tensor_dict = {} |
| 259 | states_attr = {} |
| 260 | for k, v in states.items(): |
| 261 | assert isinstance(v, tensor.Tensor), "Only tensor state is allowed" |
| 262 | tensor_dict[k] = tensor.to_numpy(v) |
| 263 | states_attr[k] = { |
| 264 | 'state_type': self.MODEL_STATE_TYPE, |
| 265 | 'shape': v.shape, |
| 266 | 'dtype': v.dtype |
| 267 | } |
| 268 | |
| 269 | for k, v in aux_states.items(): |
| 270 | assert isinstance(v, |
| 271 | tensor.Tensor), "Only tensor aux state is allowed" |
| 272 | tensor_dict[k] = tensor.to_numpy(v) |
| 273 | states_attr[k] = { |
| 274 | 'state_type': self.AUX_STATE_TYPE, |
| 275 | 'shape': v.shape, |
| 276 | 'dtype': v.dtype |
| 277 | } |
| 278 | |
| 279 | # save to files |
| 280 | timestamp = time.time() |
| 281 | tmp_dir = '/tmp/singa_save_states_%s' % timestamp |
| 282 | os.mkdir(tmp_dir) |
| 283 | tensor_dict_fp = tmp_dir + self.TENSOR_DICT_FILENAME |
| 284 | states_attr_fp = tmp_dir + self.STATES_ATTR_FILENAME |
| 285 | |
| 286 | np.savez(tensor_dict_fp, **tensor_dict) |
| 287 | |
| 288 | with open(states_attr_fp, 'w') as fp: |
| 289 | json.dump(states_attr, fp) |
| 290 | |
| 291 | compression = zipfile.ZIP_DEFLATED |
| 292 | with zipfile.ZipFile(fpath, mode="w") as zf: |
| 293 | zf.write(tensor_dict_fp, |
| 294 | os.path.basename(tensor_dict_fp), |
| 295 | compress_type=compression) |
| 296 | zf.write(states_attr_fp, |
| 297 | os.path.basename(states_attr_fp), |
| 298 | compress_type=compression) |
| 299 | |
| 300 | # clean up tmp files |
| 301 | os.remove(tensor_dict_fp) |