TensorFlow Hub：Kaggle Modelで配布されているTensorFlow v1のモデルをKerasLayerへと変換してbatch_sizeが固定されたInputと接続することにより可変値を持たないtflite形式のモデルを生成する

Kaggle Modelから入手したTensorFlowのモデルのバッチサイズを固定して書き出す

組み込み機器向けに、入力にNone(可変値)を持たない、固定されたInput Shapeを持つTensorFlow Liteモデルを作成しようと調べてみましたが、 keras.Input の batch_shape 引数を変更するだけで解決できる！ といった、なんとも呆気ない解決方法に辿り着いたため、記録しておきます。

関連記事

TensorFlow Hub：Kaggle Modelで配布されているモデルの入出力を再定義する

TensorFlow Hubで配布されている学習済みのAIモデルは入力テンソルにNone値を持つ、可変的なモデルであることが一般的です。私が開発に利用しようとしていた環境では、可変値を持つモデルを利用することができませんでした。 そのため、TensorFlow Hub：Kaggle Modelで配布されているモデルを hub.KerasLayer メソッドを利用してKerasLayerとして定義し、そのレイヤーの入力を固定的なInput Shapeとすることにより、可変値を持たないTensorFlow Liteモデルを定義しました。

### -->> FIXED Shape
input_y = keras.Input(batch_shape=(1, 1000, ), name='y')
input_z = keras.Input(batch_shape=(1, 120, ), name='z')
### -->> Dynamic Shape
#input_y = keras.Input(shape=(1000, ), name='y')
#input_z = keras.Input(shape=(120, ), name='z')

以上のshapeを持つInputとhub.KerasLayerを接続した結果のモデルの入出力は、以下のようになります。batch_sizeが1に固定されていることがわかります。

再定義されたモデルはKeras形式ですので、ONNX形式に書き出すことも可能です。

onnx_model, _ = tf2onnx.convert.from_keras(model, outputs_as_nchw=[model.outputs[0].name])
onnx.save(onnx_model, './biggan-128.onnx')

コード全文

ソースコードの全文は以下に掲載します。

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import gridspec
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow_hub as hub
import tf2onnx
import onnx
import onnxruntime
from onnxruntime.datasets import get_example

model_url = "https://www.kaggle.com/models/deepmind/biggan/TensorFlow1/128/2"
input_truncation = keras.Input(batch_shape=(), name='truncation')

### -->> FIXED Shape
input_y = keras.Input(batch_shape=(1, 1000, ), name='y')
input_z = keras.Input(batch_shape=(1, 120, ), name='z')
### -->> Dynamic Shape
#input_y = keras.Input(shape=(1000, ), name='y')
#input_z = keras.Input(shape=(120, ), name='z')

hub_layer = hub.KerasLayer(
            model_url,
            trainable=False,
            signature="default", 
            signature_outputs_as_dict=True,
        )

inputs = dict(y=input_y, z=input_z, truncation=input_truncation)
output = hub_layer(inputs)
model = tf.keras.models.Model(inputs=[input_truncation, input_y, input_z], outputs=[output])
model.build(([1], [1, 1000], [1, 120]))
model.summary()
model.save("biggan-128")
model.save("biggan-128.h5")

converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
open("biggan-128.tflite", "wb").write(tflite_model)

onnx_model, _ = tf2onnx.convert.from_keras(model, outputs_as_nchw=[model.outputs[0].name])
onnx.save(onnx_model, './biggan-128.onnx')

requirements.txt

absl-py==2.2.1
astunparse==1.6.3
cachetools==5.5.2
certifi==2025.1.31
charset-normalizer==3.4.1
coloredlogs==15.0.1
contourpy==1.3.1
cycler==0.12.1
flatbuffers==1.12
fonttools==4.57.0
gast==0.4.0
google-auth==2.38.0
google-auth-oauthlib==0.4.6
google-pasta==0.2.0
grpcio==1.71.0
h5py==3.13.0
humanfriendly==10.0
idna==3.10
keras==2.9.0
Keras-Preprocessing==1.1.2
kiwisolver==1.4.8
libclang==18.1.1
Markdown==3.7
MarkupSafe==3.0.2
matplotlib==3.10.1
mpmath==1.3.0
numpy==1.23.5
oauthlib==3.2.2
onnx==1.14.1
onnx-graphsurgeon==0.5.7
onnx2tf==1.26.9
onnxruntime==1.21.0
opt_einsum==3.4.0
packaging==24.2
pillow==11.1.0
protobuf==3.20.3
psutil==7.0.0
pyasn1==0.6.1
pyasn1_modules==0.4.2
pyparsing==3.2.3
python-dateutil==2.9.0.post0
requests==2.32.3
requests-oauthlib==2.0.0
rsa==4.9
six==1.17.0
sng4onnx==1.0.4
sympy==1.13.3
tensorboard==2.9.0
tensorboard-data-server==0.6.1
tensorflow==2.9.0
tensorflow-hub==0.16.1
tensorflow-io-gcs-filesystem==0.37.1
tensorflow-neuron==1.0
termcolor==3.0.0
tf-keras==2.14.1
tf2onnx==1.13.0
typing_extensions==4.13.0
urllib3==2.3.0
Werkzeug==3.1.3
wrapt==1.17.2

---

定義済みのKerasモデルから可変値を持たないtfliteモデルを生成する

TensorFlow Hub：Kaggle Modelからではない、自身で定義したKerasモデルの batch_size を固定する方法について、以下に掲載します。Functional APIを利用して、batch_sizeを固定したInputを作成し、これをmodelの入力とすることで解決することができます。

input = keras.layers.Input(batch_shape=(1,784), name='Input')
output = model(input)
new_model = Model(input, output)
new_model.build([28,28,1])
new_model.summary()

Mnistを分類する自身で定義したモデルのbatch_sizeを固定する

Input の batch_shape によりモデルの入出力のbatch_sizeを固定する方法は、Kaggle Modelから入手していない、自身で定義したKeras Modelに対しても有用です。以下は Mnist データセットを分類するプログラムであり、このモデルの入力である (None, 784) に対して input = keras.layers.Input(batch_shape=(1,784), name='Input') の入力を定義し、Functioanl APIを利用し new_model = Model(input, output) として接続することによって batch_size を1に固定しています。

固定後、固定する前と認識率が変わっていないことを確認しておきましょう。

### Evaluate the model
loss, acc = model.evaluate(test_images, test_labels, verbose=2)
print("Trained model, accuracy: {:5.2f}%".format(100 * acc))

######
### Connect an new input
######
input = keras.layers.Input(batch_shape=(1,784), name='Input')
output = model(input)
new_model = Model(input, output)
new_model.build([28,28,1])
new_model.summary()


### Evaluate the model
loss, acc = model.evaluate(test_images, test_labels, verbose=2)
print("Untrained new_model, accuracy: {:5.2f}%".format(100 * acc))

ソースコード全文

ソースコード全文は以下の通りです。以下の記事を参考に作成しました。

import os

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.models import Model
print(tf.version.VERSION)

### Read datasets
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
train_labels = train_labels[:1000]
test_labels = test_labels[:1000]
train_images = train_images[:1000].reshape(-1, 28 * 28) / 255.0
test_images = test_images[:1000].reshape(-1, 28 * 28) / 255.0

### Define a simple sequential model
def create_model():
  model = tf.keras.Sequential([
    keras.layers.Dense(512, activation='relu', input_shape=(784,)),
    keras.layers.Dropout(0.2),
    keras.layers.Dense(10)
  ])

  model.compile(optimizer='adam',
                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

  return model

### Create a basic model instance
model = create_model()

### Display the model's architecture
model.summary()

### Evaluate the model
loss, acc = model.evaluate(test_images, test_labels, verbose=2)
print("Untrained model, accuracy: {:5.2f}%".format(100 * acc))

### Train the model
model.fit(train_images, train_labels, epochs=5)

### Evaluate the model
loss, acc = model.evaluate(test_images, test_labels, verbose=2)
print("Trained model, accuracy: {:5.2f}%".format(100 * acc))

######
### Connect an new input
######
input = keras.layers.Input(batch_shape=(1,784), name='Input')
output = model(input)
new_model = Model(input, output)
new_model.build([28,28,1])
new_model.summary()

### Evaluate the model
loss, acc = model.evaluate(test_images, test_labels, verbose=2)
print("Untrained new_model, accuracy: {:5.2f}%".format(100 * acc))

######
### Export in tflite format
######
converter = tf.lite.TFLiteConverter.from_keras_model(new_model)
tflite_model = converter.convert()
open("dense_mnist.tflite", "wb").write(tflite_model)

ResNetのbatch_sizeを固定する場合

同じように ResNet の batch_size を固定する場合は以下のソースコードを利用できます。このソースコード中に出現する CutoffInputLayers は以下の記事を参考に定義しました。このメソッドを利用することにより、モデルの先頭の層をいくつか除去することができますので、新しい入力を定義したい場合に利用すると良いでしょう。

ソースコード全文

ソースコード全文を以下に掲載します。

import tensorflow as tf
import tensorflow.keras.layers as layers
from tensorflow.keras.applications import resnet_v2
from tensorflow.keras.utils import plot_model
from tensorflow.keras.models import Model

######
### Cut off the beginning of the Keras Layer
### -->> https://qiita.com/T-STAR/items/f3d01c06fa788fb3f635
######
def CutoffInputLayers(model, target_layer_name):
    def f(x):
        input_layers = {}
        target_layer_index = None
        for i, layer in enumerate(model.layers):
            if layer.name == target_layer_name:
                target_layer_index = i
            for node in layer._outbound_nodes:
                layer_name = node.outbound_layer.name
                if layer_name not in input_layers:
                    input_layers.update(
                            {layer_name: [layer.name]})
                else:
                    input_layers[layer_name].append(layer.name)

        if target_layer_index==None:
            raise ValueError(target_layer_name+" not found.")
        new_output_tensor = {}

        model_outputs = []
        new_output_tensor.update(
            {model.layers[target_layer_index].name: x})
        for layer in model.layers[target_layer_index+1:]:

            layer_input = [new_output_tensor[layer_aux] for layer_aux in input_layers[layer.name]]
            if len(layer_input) == 1:
                layer_input = layer_input[0]

            x = layer(layer_input)

            new_output_tensor.update({layer.name: x})
            if layer.name in model.output_names:
                model_outputs.append(x)
        return x
    return f

######
### Load a predefined model
######
resnet = resnet_v2.ResNet50V2()
resnet.summary()

######
### Remove an input and connect a new one
######
input = layers.Input(batch_shape=(1,224,224,3), name='Input')
#output = CutoffInputLayers(resnet, 'input_1')(input)
output = resnet(input)
model = Model(input, output)
model.build([224,224,3])
model.summary()

######
### Export in tflite format
######
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
open("ResNet50V2.tflite", "wb").write(tflite_model)

---

記事は以上となります。ありがとうございました。