From dc914cde50174494d2890a00a9179e655368a2ad Mon Sep 17 00:00:00 2001
From: George Hotz <george@comma.ai>
Date: Wed, 22 Feb 2023 20:08:58 -0800
Subject: [PATCH] compile_tensorflow

---
 examples/compile_efficientnet.py | 38 ++++++++++++++------------
 examples/compile_tensorflow.py   | 46 ++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 17 deletions(-)
 create mode 100644 examples/compile_tensorflow.py

diff --git a/examples/compile_efficientnet.py b/examples/compile_efficientnet.py
index 01e6ef8d80..15caa62d8f 100644
--- a/examples/compile_efficientnet.py
+++ b/examples/compile_efficientnet.py
@@ -3,23 +3,9 @@ from tinygrad.tensor import Tensor
 from extra.utils import fetch
 import ast
 
-if __name__ == "__main__":
-  model = EfficientNet(0)
-  model.load_from_pretrained()
-
-  from extra.jit import TinyJit
-  @TinyJit
-  def run(x): return model.forward(x).realize()
-
-  # twice to run the JIT
-  the_input = Tensor.randn(1,3,224,224)
-  the_output = run(the_input)
-  the_output = run(the_input)
-
-  # TODO: fetch this from the jit in self.input_replace and self.ret (hint: use get_parameters on self.ret)
-  special_names = {id(the_input.lazydata.realized.cl): "input", id(the_output.lazydata.realized.cl): "outputs"}
-
+def compile_net(run, special_names):
   # c header
+  weights = []
   cprog = ["#include <stdio.h>", "#include <math.h>","#define max(x,y) fmax(x,y)"] 
 
   # functions that run the net
@@ -49,11 +35,29 @@ if __name__ == "__main__":
   # buffers (weights)
   for name,cl in bufs_to_save.items():
     weight = ''.join(["\\x%02X"%x for x in bytes(memoryview(cl)[0:len(cl)//4])])
-    cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
+    weights.append(f"unsigned char {name}_data[] = \"{weight}\";")
     cprog.append(f"float *{name} = (float *){name}_data;")
 
   # the net
   cprog += ["void net() {"] + statements + ["}"]
+  return weights+cprog
+
+if __name__ == "__main__":
+  model = EfficientNet(0)
+  model.load_from_pretrained()
+
+  from extra.jit import TinyJit
+  @TinyJit
+  def run(x): return model.forward(x).realize()
+
+  # twice to run the JIT
+  the_input = Tensor.randn(1,3,224,224)
+  the_output = run(the_input)
+  the_output = run(the_input)
+
+  # TODO: fetch this from the jit in self.input_replace and self.ret (hint: use get_parameters on self.ret)
+  special_names = {id(the_input.lazydata.realized.cl): "input", id(the_output.lazydata.realized.cl): "outputs"}
+  cprog = compile_net(run, special_names)
 
   # image library!
   cprog += ["#define STB_IMAGE_IMPLEMENTATION", fetch("https://raw.githubusercontent.com/nothings/stb/master/stb_image.h").decode('utf-8')]
diff --git a/examples/compile_tensorflow.py b/examples/compile_tensorflow.py
new file mode 100644
index 0000000000..8c2af597c1
--- /dev/null
+++ b/examples/compile_tensorflow.py
@@ -0,0 +1,46 @@
+# An example to compile a small Tensorflow model to extremely portable C code
+
+import os
+import tensorflow as tf
+import tf2onnx
+import onnx
+from examples.compile_efficientnet import compile_net
+from extra.onnx import get_run_onnx
+from tinygrad.tensor import Tensor
+
+def get_uncompiled_model2(dataset_size=32, output_size=4):
+  inputs = tf.keras.Input(shape=(dataset_size,), name="inputs")
+  x = tf.keras.layers.Dense(16, activation="relu", name="dense_1")(inputs)
+  x = tf.keras.layers.BatchNormalization()(x)
+  x = tf.keras.layers.Dense(32, activation="relu", name="dense_2")(x)
+  outputs = tf.keras.layers.Dense(output_size, activation="sigmoid", name="predictions")(x)
+  model = tf.keras.Model(inputs=inputs, outputs=outputs)
+  return model
+
+def create_onnx_model():
+  model = get_uncompiled_model2()
+  input_signature = [tf.TensorSpec([1,32], tf.float32, name='x')]
+  onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature, opset=13)
+  return onnx_model
+
+def compile_onnx_model(onnx_model):
+  run_onnx = get_run_onnx(onnx_model)
+
+  from extra.jit import TinyJit
+  @TinyJit
+  def run(x): return run_onnx({"x": x}, debug=False)['predictions'].realize()
+
+  the_input = Tensor.randn(1,32)
+  the_output = run(the_input)
+  the_output = run(the_input)
+
+  special_names = {id(the_input.lazydata.realized.cl): "input", id(the_output.lazydata.realized.cl): "outputs"}
+  cprog = compile_net(run, special_names)
+  cprog[-1] = "return outputs;\n}"
+
+  print('\n'.join(cprog).replace("void net()", "float *infer(float *input)").replace("float input[32];\n", ""))
+
+if __name__ == "__main__":
+  onnx_model = create_onnx_model()
+  compile_onnx_model(onnx_model)
+