keras-tensorflow-onnx notes

.gitignore

keras_model.h5
tf_model
model.onnx
updated_tf_model
updated_keras_model.h5

README.md

objective

Convert Keras models with data-format channels_first (NCHW) to channels_last (NHWC).

instructions

• setup environment

docker run -it -v $PWD:/workdir -w /workdir tensorflow/tensorflow:2.8.0 bash
pip install google==3.0.0 tf2onnx==1.12.0 onnx-tf==1.10.0 tensorflow_probability==0.15.0 SimpleITK==2.1.1.2 scikit-image==0.19.3

• copy keras model to /workdir, assuming model filename to be keras_model.h5.

• run below to convert to tf model.

python keras_to_tf.py keras_model.h5 tf_model

• convert tf to onnx model file.

python -m tf2onnx.convert --saved-model tf_model --output model.onnx

• the above 2 steps could be replaced by below - keras model direct to onnx model - if below doesn't error out for you

python -m tf2onnx.convert --keras keras_model.h5 --output model.onnx

• convert onnx back to tf model format.

onnx-tf convert --infile model.onnx --outdir updated_tf_model

• sample test inference *** with hard coded input shape and tensor names ***

python tf_inference.py updated_tf_model

####### *** TODO: below not working ***

• convert keras to onnx then back to keras

apt-get update;apt-get upgrade -yq;
vim /usr/local/lib/python3.8/dist-packages/onnx2keras/activation_layers.py

change line 144 to 151 per suggested by  
https://github.com/gmalivenko/onnx2keras/pull/107/files

+    softmax_layer = keras.layers.Softmax(axis=params.get('axis',-1), name=keras_name)
+    layers[node_name] = softmax_layer(input_0)
+    layers[node_name].set_shape(layers[node_name].shape)

-    def target_layer(x, axis=params['axis']):
-        import tensorflow as tf
-        return tf.nn.softmax(x, axis=axis)
-
-    lambda_layer = keras.layers.Lambda(target_layer, name=keras_name)
-    layers[node_name] = lambda_layer(input_0)
-    layers[node_name].set_shape(layers[node_name].shape)
-    lambda_func[keras_name] = target_layer

python keras_to_onnx_to_keras.py keras_model.h5 updated_keras_model.h5
python keras_inference.py updated_keras_model.h5

references

    
https://www.tensorflow.org/api_docs/python/tf/saved_model/load
https://github.com/onnx/tensorflow-onnx
https://github.com/onnx/onnx-tensorflow/blob/main/doc/CLI.md
https://github.com/onnx/tensorflow-onnx/blob/e896723e410a59a600d1a73657f9965a3cbf2c3b/tf2onnx/convert.py#L408
https://stackoverflow.com/questions/37689423/convert-between-nhwc-and-nchw-in-tensorflow
https://stackoverflow.com/questions/63682625/how-to-convert-saved-model-from-nchw-to-nhwc/63867165
https://gist.github.com/MartinNowak/9a2e467d740c43fe7e6fef6c09502449
https://github.com/onnx/onnx-tensorflow/issues/862


pip install onnx2keras==0.0.24 keras2onnx==1.7.0

keras_inference.py

'''from onnx-tf output
2022-08-10 00:07:11,034 - INFO - Model inputs: ['conv2d_5_input']
2022-08-10 00:07:11,034 - INFO - Model outputs: ['dense_6']
'''
import sys
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model

my_input_shape = (20,1,256,256)
batch_data = np.random.rand(*my_input_shape).astype(np.float32)

model_folder = sys.argv[1]
model = load_model(model_folder)
output = model.predict(batch_data,batch_size=8)
print(output.shape)
print(np.argmax(output,axis=-1))

keras_to_onnx_to_keras.py

import sys
import onnx

from tf2onnx import convert
from onnx2keras import onnx_to_keras
from tensorflow.keras.models import load_model

keras_model_file = sys.argv[1]
updated_keras_model_file = sys.argv[2]

keras_model = load_model(keras_model_file)
print(keras_model)
model_proto, external_tensor_storage = convert.from_keras(keras_model)
print(model_proto)
input_names=None
input_shapes=None
name_policy='short'
verbose=True
change_ordering=True

keras_model = onnx_to_keras(
    onnx_model=model_proto, input_names=input_names,
    input_shapes=input_shapes, name_policy=name_policy,
    verbose=verbose, change_ordering=change_ordering)
keras_model.summary()
keras_model.save(updated_keras_model_file)

keras_to_tf.py

import os
import sys
from tensorflow.keras.models import load_model

keras_model_file = sys.argv[1]
tf_model_folder = sys.argv[2]

model = load_model(keras_model_file)

if 'channels_first' in model.to_json():
    pass
else:
    raise ValueError("channels_first phrase not found, no need to convert!")

model.summary()
model.save(tf_model_folder)

json_file = os.path.join(tf_model_folder,f'{keras_model_file}.json')
with open(json_file,'w') as f:
    f.write(model.to_json())

'''
import google.protobuf
from tensorflow.core.protobuf import saved_model_pb2
import tensorflow as tf
# ???
with open(os.path.join(tf_model_folder, 'keras_metadata.pb'), 'rb') as f:
    message = f.read()
with open(os.path.join(tf_model_folder, 'keras_metadata.pb.txt'), 'w') as f:
    f.write(google.protobuf.text_format.MessageToString(message))
'''

'''
docker run -it -v $PWD:/workdir -w /workdir tensorflow/tensorflow:2.8.0-gpu-jupyter bash

docker run -it -v $PWD:/workdir -w /workdir tensorflow/tensorflow:2.8.0 bash
pip install google==3.0.0 tf2onnx==1.12.0 onnx-tf==1.10.0 tensorflow_probability==0.15.0 SimpleITK==2.1.1.2 scikit-image==0.19.3
# mv model file to folder rename as `keras_model.h5`
# convert keras to tf
python keras-to-tf.py keras_model.h5 tf_model
# convert tf to onnx
python -m tf2onnx.convert --saved-model tf_model --output model.onnx
# convert onnx back to tf
onnx-tf convert --infile model.onnx --outdir updated_tf_model

cp tf_model/keras_metadata.pb updated_tf_model

python tf-to-keras.py updated_tf_model updated_keras.h5

'''

tf_inference.py

'''from onnx-tf output
2022-08-10 00:07:11,034 - INFO - Model inputs: ['conv2d_5_input']
2022-08-10 00:07:11,034 - INFO - Model outputs: ['dense_6']
'''
import sys
import numpy as np
import tensorflow as tf

my_input_shape = (20,1,256,256)
input_tensor_name = 'input_1'
output_tensor_name ='dense_3'
batch_data = np.random.rand(*my_input_shape).astype(np.float32)
batch_data = tf.constant(batch_data,name=input_tensor_name)

model_folder = sys.argv[1]
loaded = tf.saved_model.load(model_folder)
#print(loaded)
model_fn = loaded.signatures['serving_default']
#print(model_fn)
#print(dir(model_fn))
#print(model_fn.graph)

output_dict = model_fn(batch_data)
output0  = output_dict[output_tensor_name]
print(output0.shape)
#print(np.argmax(output,axis=-1))

# model_file = sys.argv[2]
# model = tf.keras.models.load_model(model_file)
# output1 = model.predict(batch_data)
# print(output1.shape)
# for x,y in zip(output0,output1):
#     print(x,y)

tf_to_keras.py

import numpy as np
import sys    
import tensorflow as tf

tf_model_file = sys.argv[1]
keras_model_file = sys.argv[2]

model = tf.keras.models.load_model(tf_model_file)

# loaded = tf.saved_model.load(tf_model_file)
# print(dir(loaded))
# model_fn = loaded.signatures['serving_default']
# print(dir(model_fn))



#tf.keras.models.save_model(model,keras_model_file, save_format='h5')

# with open(keras_json_file,'r') as f:
#     keras_json=f.read()
# loaded_model = tf.keras.models.model_from_json(keras_json)

#_ = loaded_model.load_weights(tf_model_file)
#batch_data = np.random.rand(1,256,256,1).astype(np.float32)
#out = loaded_model.predict(batch_data)
#print(out.shape)

# loaded = tf.saved_model.load(tf_model_file)
# print(dir(loaded))
# model_fn = loaded.signatures['serving_default']
# print(dir(model_fn))

provided my answer to relevant question in SO:
https://stackoverflow.com/a/73311512/868736

  # fugly custom batching
  out_list = []
  img_split = np.array_split(img,np.ceil(img.shape[0]/batch_size),axis=0)
  for item in img_split:
      batch_data = tf.constant(item,name=input_tensor_name)
      output_dict = model_fn(batch_data)
      out = output_dict[output_tensor_name]
      out = out.numpy()
      out_list.append(out)
  out = np.concatenate(out_list,axis=0)