This code is generated automatically using "go generate ./..." from model/tfkg_model.py. DO NOT EDIT manually.
(customDefinitions []string)
| 6 | |
| 7 | // This code is generated automatically using "go generate ./..." from model/tfkg_model.py. DO NOT EDIT manually. |
| 8 | func GetTfkgPythonCode(customDefinitions []string) string { |
| 9 | return strings.ReplaceAll(`import json |
| 10 | import os |
| 11 | import logging |
| 12 | import sys |
| 13 | |
| 14 | import tensorflow as tf |
| 15 | import numpy as np |
| 16 | |
| 17 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # ERROR |
| 18 | logging.getLogger('tensorflow').setLevel(logging.ERROR) |
| 19 | logging.disable(logging.WARNING) |
| 20 | |
| 21 | custom_objects = {} |
| 22 | |
| 23 | # tfkg-custom-definitions |
| 24 | |
| 25 | with open(sys.argv[1], "r") as f: |
| 26 | config = json.load(f) |
| 27 | |
| 28 | def save_model(dir): |
| 29 | model = tf.keras.models.model_from_json(config["model_config"], custom_objects=custom_objects) |
| 30 | |
| 31 | weights_spec = [] |
| 32 | for item in model.weights: |
| 33 | weights_spec.append(tf.TensorSpec(shape=item.shape, dtype=item.dtype)) |
| 34 | |
| 35 | if config["model_definition_save_dir"] != "": |
| 36 | summary = [] |
| 37 | model.summary(print_fn=lambda x: summary.append(x)) |
| 38 | with open(config["model_definition_save_dir"] + "/model-summary.txt", "w") as f: |
| 39 | f.write("\n".join(summary)) |
| 40 | weight_names = [] |
| 41 | for item in model.weights: |
| 42 | weight_names.append(item.name) |
| 43 | with open(config["model_definition_save_dir"] + "/weight_names.json", "w") as f: |
| 44 | json.dump(weight_names, f) |
| 45 | |
| 46 | learn_signature = [] |
| 47 | predict_input_signature = [] |
| 48 | |
| 49 | zero_inputs = [] |
| 50 | |
| 51 | model_config = json.loads(config["model_config"]) |
| 52 | |
| 53 | for model_layer in model_config["config"]["layers"]: |
| 54 | if model_layer["class_name"] == "InputLayer": |
| 55 | input_shape = [config["batch_size"]] |
| 56 | predict_input_shape = [None] |
| 57 | for dim in model_layer["config"]["batch_input_shape"][1:]: |
| 58 | input_shape.append(dim) |
| 59 | predict_input_shape.append(dim) |
| 60 | zero_inputs.append( |
| 61 | tf.zeros(shape=input_shape, dtype=model_layer["config"]["dtype"]) |
| 62 | ) |
| 63 | learn_signature.append(tf.TensorSpec( |
| 64 | shape=input_shape, |
| 65 | dtype=model_layer["config"]["dtype"], |