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)