Method save_states

examples/cnn_ms/pkg_model_code/model.py:247–306

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={})

            return self.forward(*input, **kwargs)

    def save_states(self, fpath, aux_states={}):
        """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
        """
        assert not os.path.isfile(fpath), (
            "Failed to save states, %s is already existed." % fpath)

        states = self.get_states()

        # save states data and attr
        tensor_dict = {}
        states_attr = {}
        for k, v in states.items():
            assert isinstance(v, tensor.Tensor), "Only tensor state is allowed"
            tensor_dict[k] = tensor.to_numpy(v)
            states_attr[k] = {
                'state_type': self.MODEL_STATE_TYPE,
                'shape': v.shape,
                'dtype': v.dtype
            }

        for k, v in aux_states.items():
            assert isinstance(v,
                              tensor.Tensor), "Only tensor aux state is allowed"
            tensor_dict[k] = tensor.to_numpy(v)
            states_attr[k] = {
                'state_type': self.AUX_STATE_TYPE,
                'shape': v.shape,
                'dtype': v.dtype
            }

        # save to files
        timestamp = time.time()
        tmp_dir = '/tmp/singa_save_states_%s' % timestamp
        os.mkdir(tmp_dir)
        tensor_dict_fp = tmp_dir + self.TENSOR_DICT_FILENAME
        states_attr_fp = tmp_dir + self.STATES_ATTR_FILENAME

        np.savez(tensor_dict_fp, **tensor_dict)

        with open(states_attr_fp, 'w') as fp:
            json.dump(states_attr, fp)

        compression = zipfile.ZIP_DEFLATED
        with zipfile.ZipFile(fpath, mode="w") as zf:
            zf.write(tensor_dict_fp,
                     os.path.basename(tensor_dict_fp),
                     compress_type=compression)
            zf.write(states_attr_fp,
                     os.path.basename(states_attr_fp),
                     compress_type=compression)

        # clean up tmp files
        os.remove(tensor_dict_fp)