def save(self, export_dir: Text, options: Optional[tf.saved_model.SaveOptions] = None): """Save the policy to the given `export_dir`. Args: export_dir: Directory to save the policy to. options: Optional `tf.saved_model.SaveOptions` object. """ tf.compat.v2.saved_model.save(self._policy, export_dir, signatures=self._signatures, options=options) temp_spec_file_name = '{}_temp'.format(POLICY_SPECS_PBTXT) temp_spec_output_path = os.path.join(export_dir, temp_spec_file_name) specs = { 'collect_data_spec': self._policy.collect_data_spec, 'policy_state_spec': self._policy.policy_state_spec } tensor_spec.to_pbtxt_file(temp_spec_output_path, specs) spec_output_path = os.path.join(export_dir, POLICY_SPECS_PBTXT) # By moving the file to its final location makes it safer to wait for the # file (e.g. from a separate binary). The parameter `overwrite=True` # reproduces the exact previous behavior. tf.io.gfile.rename(temp_spec_output_path, spec_output_path, overwrite=True)
def save(self, export_dir): """Save the policy to the given `export_dir`.""" tf.saved_model.save(self._policy, export_dir, signatures=self._signatures) spec_output_path = os.path.join(export_dir, COLLECT_POLICY_SPEC) specs = { 'collect_data_spec': self._policy.collect_data_spec, 'policy_state_spec': self._policy.policy_state_spec } tensor_spec.to_pbtxt_file(spec_output_path, specs)
def save(self, export_dir, options=None): """Save the policy to the given `export_dir`. Args: export_dir: Directory to save the policy to. options: Optional `tf.saved_model.SaveOptions` object. """ tf.compat.v2.saved_model.save(self._policy, export_dir, signatures=self._signatures, options=options) spec_output_path = os.path.join(export_dir, POLICY_SPECS_PBTXT) specs = { 'collect_data_spec': self._policy.collect_data_spec, 'policy_state_spec': self._policy.policy_state_spec } tensor_spec.to_pbtxt_file(spec_output_path, specs)