我有一個通過模型子類化撰寫的 tensorflow 模型,我想將其匯出為 ONNX 格式。這很簡單,附帶腳本。但是,ONNX 模型的輸入變數的名稱是args_0. 我該如何重命名它?
import tensorflow as tf
import tf2onnx
from tensorflow.python.keras import Model
from tensorflow.python.keras.layers import Dense
class MyModel(Model):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.dense = Dense(16)
def call(self, inputs, **kwargs):
return self.dense(inputs)
def to_onnx(self, output_path, opset=14):
model_proto, _ = tf2onnx.convert.from_keras(
self,
input_signature=[tf.TensorSpec((1, 128))],
opset=opset,
output_path=output_path,
)
return
if __name__ == "__main__":
output_path = "./test.onnx"
A = MyModel()
A.to_onnx(output_path)
uj5u.com熱心網友回復:
input_signature您可以像這樣提供輸入名稱name="input_name",它應該看起來像
input_signature=[tf.TensorSpec((1, 128), name="input_name")],
如本筆記本所示
轉載請註明出處,本文鏈接:https://www.uj5u.com/shujuku/414254.html
標籤:
