Attempt to try tf2

Author: vanpelt

examples/keras-perf/cnn.py

@@ -1,13 +1,7 @@
-from keras import backend as K
-from keras.datasets import mnist
-from keras.models import Sequential
-from keras.layers import Conv2D, MaxPooling2D, Dropout, Dense, Flatten
-from keras.utils import np_utils
-from wandb.keras import WandbCallback
-from keras.callbacks import TensorBoard
-import wandb
 import os
 import tensorflow as tf
+import wandb
+from wandb.keras import WandbCallback
 
 run = wandb.init()
 config = run.config
@@ -20,7 +14,7 @@
 config.img_height = 28
 config.epochs = 4
 
-(X_train, y_train), (X_test, y_test) = mnist.load_data()
+(X_train, y_train), (X_test, y_test) = tf.keras.datasets.mnist.load_data()
 
 X_train = X_train.astype('float32')
 X_train /= 255.
@@ -34,68 +28,35 @@
     X_test.shape[0], config.img_width, config.img_height, 1)
 
 # one hot encode outputs
-y_train = np_utils.to_categorical(y_train)
-y_test = np_utils.to_categorical(y_test)
+y_train = tf.keras.utils.to_categorical(y_train)
+y_test = tf.keras.utils.to_categorical(y_test)
 num_classes = y_test.shape[1]
 labels = range(10)
 
 # build model
-model = Sequential()
-model.add(Conv2D(32,
-                 (config.first_layer_conv_width, config.first_layer_conv_height),
-                 input_shape=(28, 28, 1),
-                 activation='relu'))
-model.add(MaxPooling2D(pool_size=(2, 2)))
-model.add(Flatten())
-model.add(Dense(config.dense_layer_size, activation='relu'))
-model.add(Dense(num_classes, activation='softmax'))
-
+model = tf.keras.models.Sequential()
+model.add(tf.keras.layers.Conv2D(32,
+                                 (config.first_layer_conv_width,
+                                  config.first_layer_conv_height),
+                                 input_shape=(28, 28, 1),
+                                 activation='relu'))
+model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
+model.add(tf.keras.layers.Flatten())
+model.add(tf.keras.layers.Dense(config.dense_layer_size, activation='relu'))
+model.add(tf.keras.layers.Dense(num_classes, activation='softmax'))
 model.compile(loss='categorical_crossentropy', optimizer='adam',
-              metrics=['accuracy'], weighted_metrics=['accuracy'])
-
-
+              metrics=['accuracy'])
+# log the number of total parameters
+config.total_params = model.count_params()
+print("Total params: ", config.total_params)
 model.fit(X_train, y_train, validation_data=(X_test, y_test),
           epochs=config.epochs,
           callbacks=[WandbCallback(data_type="image", save_model=False),
-                     TensorBoard(log_dir=wandb.run.dir)])
-
-
-def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
-    """
-    Freezes the state of a session into a pruned computation graph.
-
-    Creates a new computation graph where variable nodes are replaced by
-    constants taking their current value in the session. The new graph will be
-    pruned so subgraphs that are not necessary to compute the requested
-    outputs are removed.
-    @param session The TensorFlow session to be frozen.
-    @param keep_var_names A list of variable names that should not be frozen,
-                          or None to freeze all the variables in the graph.
-    @param output_names Names of the relevant graph outputs.
-    @param clear_devices Remove the device directives from the graph for better portability.
-    @return The frozen graph definition.
-    """
-    graph = session.graph
-    with graph.as_default():
-        freeze_var_names = list(
-            set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
-        output_names = output_names or []
-        output_names += [v.op.name for v in tf.global_variables()]
-        input_graph_def = graph.as_graph_def()
-        if clear_devices:
-            for node in input_graph_def.node:
-                node.device = ""
-        frozen_graph = tf.graph_util.convert_variables_to_constants(
-            session, input_graph_def, output_names, freeze_var_names)
-        return frozen_graph
-
-
-model.save("cnn.h5")
-saver = tf.train.Saver()
-sess = K.get_session()
-saver.save(sess, './keras_model.ckpt')
-tf.train.write_graph(sess.graph_def, '.', 'keras_model.pbtxt')
-
-frozen_graph = freeze_session(K.get_session(), output_names=[
-                              out.op.name for out in model.outputs])
-tf.train.write_graph(frozen_graph, ".", "cnn.pb", as_text=False)
+                     tf.keras.callbacks.TensorBoard(profile_batch=3)])
+model.save('cnn.h5')
+
+# Convert to TensorFlow Lite model.
+converter = tf.lite.TFLiteConverter.from_keras_model_file('cnn.h5')
+converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
+tflite_model = converter.convert()
+open("cnn.tflite", "wb").write(tflite_model)

examples/keras-perf/train.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+virtualenv --no-site-packages venv
+. venv/bin/activate
+pip install tensorflow-gpu==2.0.0b1 wandb pillow
+
+python cnn.py
+
+deactivate