File size: 1,698 Bytes
d7755fe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | 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()
|