Tensorflow support

[attila-dusnoki-htec]

Tracking issue for all tensorflow related tasks.

[attila-dusnoki-htec]

With the release of Tensorflow 2, they deprecated freeze_graph​
And introduced a new format called SavedModel​. This can be easily shared and further tuned if needed, since the weights are stored separatelly.

Model zoos with tensorflow: tfhub.dev, kaggle​

To load this model in c++, they added a new API.
But there is no prebuilt library provided. Only for C. (which has an (3+ years old) experimental api for this)

It requires manually building tensorflow. This can be done with bazel using the //tensorflow:libtensorflow_cc.so target.
To get a proper rocm build, there is a dedicated repository for it.

[attila-dusnoki-htec]

Instead of depending on TF, we should create a tf2 parse in migraphx which can handle the new format.

[attila-dusnoki-htec]

The SavedModel will contain the following format: (source)

assets/
assets.extra/
variables/
    variables.data-?????-of-?????
    variables.index
saved_model.pb

An example SavedModel: https://www.kaggle.com/models/tensorflow/resnet-50 (It only has variables, no assets)

Those names come from here: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/cc/saved_model/constants.h
And the variables naming: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/util/tensor_bundle/naming.cc

The saved_model proto is here
These will be used to generate the c++ code, and parse the binaries with these.

Inspecting any protobuf:
pip install protobuf-inspector ansi2txt
protobuf_inspector < saved_model.pb | ansi2txt &> saved_model.dump
A partial config file to get better logs. See CONFIG.md.

Some of these proto-s are already in MGX, and the matching protos in TF and TSL.
We will need to update these as well, since never versions available.
Or go with "tf2" folder and write it from scratch. But it would be ideal to have a "common" set.

Parsing the protobuf will contain a meta_graph, which has the following: GraphDef, SaverDef, CollectionDef, SignatureDef, AssetFileDef, SavedObjectGraph.
The GraphDef is what currently parsed as tf1.
The SavedObjectGraph will hold references to the actual variables (Check for "VARIABLE_VALUE" in pb dump). Each path is in variables.index file, which probably holds the weight location in variables.data-*.

Here are some files to check for start: loader.cc, reader.cc, env.cc

[attila-dusnoki-htec]

Check https://github.com/onnx/tensorflow-onnx to see if it can be used to avoid/workaround tf2 implementation.

[attila-dusnoki-htec]

Tensorflow -> ONNX

Official documentation is here.

Setup

python3 -m venv tf_onnx
. tf_onnx/bin/activate

pip install tensorflow tf2onnx

Test with ResNet50

curl -L -o ./resnet50-model.tar.gz\
  https://www.kaggle.com/api/v1/models/tensorflow/resnet-50/tensorFlow2/classification/1/download
mkdir resnet50
tar xzvf resnet50-model.tar.gz -C resnet50
python -m tf2onnx.convert --saved-model resnet50 --output resnet50.onnx

It should compile fine.

Test with BERT

curl -L -o ./bert-model.tar.gz\
  https://www.kaggle.com/api/v1/models/google/bert/tensorFlow2/answer-equivalence-bem/1/download
mkdir bert
tar xzvf bert-model.tar.gz -C bert
python -m tf2onnx.convert --saved-model bert --output bert.onnx

It will fail, because Xla* operators are not supported.
see here and here

curl -L -o ./bert-qa-model.tar.gz\
  https://www.kaggle.com/api/v1/models/seesee/bert/tensorFlow2/uncased-tf2-qa/1/download
mkdir bert-qa
tar xzvf bert-model.tar.gz -C bert-qa
python -m tf2onnx.convert --saved-model bert-qa --output bert-qa.onnx

Sadly similar results :(

Test with SDXL

No premade model available, the closest is:
TODO: https://www.tensorflow.org/tutorials/generative/generate_images_with_stable_diffusion
Note: Make sure to disable JIT

[attila-dusnoki-htec]

What is uploaded to HuggingFace, is either in safetensor or h5. No SavedModel format.

In theory, all it would take to load convert it is:

import tf2onnx
import onnx
from tensorflow.keras.models import load_model

model = load_model('path/to/keras_model_you_want_to_convert.h5')

onnx_model, _ = tf2onnx.convert.from_keras(model)
onnx.save(onnx_model, 'new_model.onnx')

It might work for some, but in practice, sadly this is not the case. Because the model config is not saved, only the weights.
There are classes to help in transformers e.g. TFBertModel, which knows the config. But in that case, optimum would already do a better job with converting.