Converting KataGo to Core ML sample code

katago_to_coreml.py

# Append the following code to katago/python/play.py

import tempfile
import os
from tensorflow.python.tools.freeze_graph import freeze_graph
import coremltools as ct

model_dir = tempfile.mkdtemp()
graph_def_file = os.path.join(model_dir, 'tf_graph.pb')
checkpoint_file = os.path.join(model_dir, 'tf_model.ckpt')
frozen_graph_file = os.path.join(model_dir, 'KataGoModel.pb')
mlmodel_file = frozen_graph_file.replace("pb","mlpackage")

output_names = [
  model.policy_output.op.name,
  model.value_output.op.name,
  model.ownership_output.op.name,
  model.miscvalues_output.op.name,
  model.moremiscvalues_output.op.name
]

print(output_names)

# session_config = tf.compat.v1.ConfigProto(allow_soft_placement=True)
# session_config.gpu_options.per_process_gpu_memory_fraction = 0.3
with tf.compat.v1.Session() as session:
  saver.restore(session, model_variables_prefix)

  tf.train.write_graph(session.graph, model_dir, graph_def_file, as_text=False)
  # save the weights
  saver = tf.train.Saver()
  saver.save(session, checkpoint_file)

  # take the graph definition and weights
  # and freeze into a single .pb frozen graph file
  freeze_graph(input_graph=graph_def_file,
               input_saver="",
               input_binary=True,
               input_checkpoint=checkpoint_file,
               output_node_names=','.join(output_names),
               restore_op_name="save/restore_all",
               filename_tensor_name="save/Const:0",
               output_graph=frozen_graph_file,
               clear_devices=True,
               initializer_nodes="")

  mlmodel = ct.convert(frozen_graph_file, convert_to="mlprogram")
  mlmodel.save(mlmodel_file)

  print("Core ML model saved at {}".format(mlmodel_file))

# Run the following command
# wget https://media.katagotraining.org/uploaded/networks/zips/kata1/kata1-b40c256-s11840935168-d2898845681.zip
# unzip kata1-b40c256-s11840935168-d2898845681.zip
# python python/play.py -saved-model-dir kata1-b40c256-s11840935168-d2898845681/saved_model -name-scope swa_model