| from pathlib import Path |
| import os |
| import shutil |
|
|
| os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") |
| os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2") |
|
|
| import tensorflow as tf |
|
|
|
|
| LAB_DIR = Path(__file__).resolve().parent |
| MODEL_DIR = LAB_DIR / "savedmodel_savev2_default_write_variant" |
| OUTPUT_PREFIX_NAME = "savev2_runtime_ckpt" |
| OUTPUT_PREFIX = LAB_DIR / OUTPUT_PREFIX_NAME |
| MARKER = "TFSM_SAVEV2_DEFAULT_WRITE_MARKER_2026" |
|
|
|
|
| def remove_outputs() -> None: |
| for path in LAB_DIR.glob(f"{OUTPUT_PREFIX_NAME}*"): |
| if path.is_dir(): |
| shutil.rmtree(path) |
| else: |
| path.unlink(missing_ok=True) |
|
|
|
|
| class SaveV2DefaultWrite(tf.Module): |
| def __init__(self) -> None: |
| super().__init__() |
| self.output_prefix = tf.constant(OUTPUT_PREFIX_NAME) |
|
|
| @tf.function(input_signature=[]) |
| def fixed_savev2_write(self): |
| save_op = tf.raw_ops.SaveV2( |
| prefix=self.output_prefix, |
| tensor_names=tf.constant(["marker"], tf.string), |
| shape_and_slices=tf.constant([""], tf.string), |
| tensors=[tf.constant(MARKER)], |
| ) |
| with tf.control_dependencies([save_op] if save_op is not None else []): |
| return {"done": tf.constant(True)} |
|
|
|
|
| def main() -> None: |
| if MODEL_DIR.exists(): |
| shutil.rmtree(MODEL_DIR) |
| remove_outputs() |
|
|
| module = SaveV2DefaultWrite() |
| tf.saved_model.save( |
| module, |
| str(MODEL_DIR), |
| signatures={"serving_default": module.fixed_savev2_write}, |
| ) |
|
|
| print(f"tensorflow={tf.__version__}") |
| print(f"model_dir={MODEL_DIR}") |
| print(f"output_prefix={OUTPUT_PREFIX}") |
| print(f"marker={MARKER}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
