compile_tensorflow: add initialize and tests

examples/compile_efficientnet.py

@@ -5,7 +5,6 @@ import ast
 
 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
@@ -29,18 +28,7 @@ def compile_net(run, special_names):
       cargs.append(bufs[key][0])
     statements.append(f"{fxn.clprg.name}({', '.join(cargs)});")
 
-  # buffers (empty)
-  cprog += [f"float {x[0]}[{x[1]}];" for x in bufs.values() if x[0] not in bufs_to_save] 
-
-  # 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])])
-    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
+  return cprog, statements, bufs, bufs_to_save
 
 if __name__ == "__main__":
   model = EfficientNet(0)
@@ -57,7 +45,20 @@ if __name__ == "__main__":
 
   # 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)
+
+  cprog, statements, bufs, bufs_to_save = compile_net(run, special_names)
+
+  # buffers (empty)
+  cprog += [f"float {x[0]}[{x[1]}];" for x in bufs.values() if x[0] not in bufs_to_save] 
+
+  # 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}\";")
+    cprog.append(f"float *{name} = (float *){name}_data;")
+
+  # the net
+  cprog += ["void net() {"] + statements + ["}"]
 
   # image library!
   cprog += ["#define STB_IMAGE_IMPLEMENTATION", fetch("https://raw.githubusercontent.com/nothings/stb/master/stb_image.h").decode('utf-8')]

examples/compile_tensorflow.py

@@ -1,9 +1,13 @@
 # An example to compile a small Tensorflow model to extremely portable C code
 
-import os
+import os, sys
+os.environ["CLANG"] = '1'
+os.environ["GPU"] = '1'
+
+import numpy as np
+import subprocess
 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
@@ -17,10 +21,9 @@ def get_uncompiled_model2(dataset_size=32, output_size=4):
   model = tf.keras.Model(inputs=inputs, outputs=outputs)
   return model
 
-def create_onnx_model():
-  model = get_uncompiled_model2()
+def create_onnx_model(keras_model):
   input_signature = [tf.TensorSpec([1,32], tf.float32, name='x')]
-  onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature, opset=13)
+  onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=13)
   return onnx_model
 
 def compile_onnx_model(onnx_model):
@@ -35,12 +38,59 @@ def compile_onnx_model(onnx_model):
   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}"
+  cprog, statements, bufs, bufs_to_save = compile_net(run, special_names)
+  cprog = ["#include <string.h>", "#include <stdio.h>"] + cprog
 
-  print('\n'.join(cprog).replace("void net()", "float *infer(float *input)").replace("float input[32];\n", ""))
+  # buffers (all except input)
+  cprog += [f"float {x[0]}[{x[1]}];" for x in bufs.values() if x[0] != "input"]
+
+  # weights
+  cprog.append("void initialize(float *weights) {")
+  weights = bytes()
+  for name,cl in bufs_to_save.items():
+    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {len(cl)});")
+    weights += bytes(memoryview(cl)[0:len(cl)//4])
+  cprog.append("}")
+
+  # the net
+  cprog += ["float *infer(float *input) {"] + statements + ["return outputs;", "}"]
+
+  # test program
+  cprog.append("""int main(int argc, char *argv[]) {
+    float input[32];
+    for (int i = 0; i < 32; i++) scanf("%f", &input[i]);
+    initialize((float *)weights);
+    float *outputs = infer(input);
+    printf("%f %f %f %f\\n", outputs[0], outputs[1], outputs[2], outputs[3]);
+  }""")
+
+  # the (test) weights
+  joined_weights = ''.join(['\\x%02X'%x for x in weights])
+  cweights = f"unsigned char weights[] = \"{joined_weights}\";\n"
+
+  # ready the program
+  prg = '\n'.join(cprog)
+  print(prg)
+
+  # add test weights
+  prg = cweights + prg
+  subprocess.check_output(['clang', '-O2', '-lm', '-fPIC', '-x', 'c', '-', '-o', "/tmp/test"], input=prg.encode('utf-8'))
+
+  tinygrad_output = [x for x in the_output.numpy()[0]]
+  print("tinygrad:", tinygrad_output, file=sys.stderr)
+
+  c_input = ' '.join(["%f" % x for x in the_input[0].numpy()])+"\n"
+  c_output = [float(x) for x in subprocess.check_output(["/tmp/test"], input=c_input.encode('utf-8')).decode('utf-8').strip().split(" ")]
+  print("compiled:", c_output, file=sys.stderr)
+
+  np.testing.assert_allclose(tinygrad_output, c_output, atol=1e-5, rtol=1e-5)
+  return the_input.numpy(), c_output
 
 if __name__ == "__main__":
-  onnx_model = create_onnx_model()
-  compile_onnx_model(onnx_model)
+  keras_model = get_uncompiled_model2()
+  onnx_model = create_onnx_model(keras_model)
+  test_input, test_output = compile_onnx_model(onnx_model)
+  tf_output = keras_model(test_input).numpy()[0]
+  print("keras:   ", tf_output, file=sys.stderr)
+  np.testing.assert_allclose(tf_output, test_output, atol=1e-5, rtol=1e-5)