TensorFlow2模型转ONNX Runtime运行失败，自定义层索引问题求助

问题概述 #

将TensorFlow2模型部署到ONNX Runtime时，模型末尾的自定义计算逻辑导致运行报错。移除该逻辑（代码中here!至to here区间）改用sigmoid输出时，流程可正常运行，推测问题出在自定义层的变量索引方式上。

相关代码 #

模型定义与转ONNX代码 #

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Dropout, Concatenate, Lambda
from tensorflow.keras.models import Model
import tf2onnx

fd = np.random.rand(100,16)
fl = np.random.randint(0,1,[100,])
D = 0.4375

def get_m():
    I = Input(shape=(None, 16))
    d1 = Dense(16, kernel_regularizer='l2', activation='tanh')(I)
    dr1 = Dropout(D)(d1)
    d2 = Dense(16, kernel_regularizer='l2', activation='tanh')(dr1)
    dr2 = Dropout(D)(d2)
    d3 = Dense(16, kernel_regularizer='l2', activation='tanh')(dr2)
    dr3_1 = Dropout(D)(d3)
    dr3_2 = Dropout(D)(d3)
    d4_1 = Dense(1, kernel_regularizer='l2', activation='relu')(dr3_1)
    d4_2 =  Dense(1, kernel_regularizer='l2', activation='relu')(dr3_2)
    # here!
    c = Concatenate()([d4_1, d4_2, I])
    def custom_layer(x):
        return tf.exp(-x[:, 4] / (1e-9 + x[:, 0]) - x[:, 5] / (1e-9 + x[:, 1]))

    hl = tf.keras.layers.Lambda(custom_layer, output_shape=(1,))(c)
    #d5 = Dense(1, activation='sigmoid')(c)
    # to here
    return Model(inputs=I, outputs=hl)

model = get_m()
model.compile(loss=tf.keras.losses.BinaryFocalCrossentropy(), 
              optimizer=tf.keras.optimizers.Adam(), 
              metrics=[tf.keras.metrics.AUC(), tf.keras.metrics.Precision(), tf.keras.metrics.Recall()])
model.fit(x=fd, y=fl, epochs=100, validation_split=0.1)
tf2onnx.convert.from_keras(model, output_file='tf2_neural_model_2deep.onnx')

ONNX Runtime运行代码 #

const ort = require('onnxruntime-node');

async function runInference() {
    const sess = await ort.InferenceSession.create('tf2_neural_model_2deep.onnx');
    const rowVectors = [[0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.]];
    const inputTensor = new ort.Tensor('float32', rowVectors.flat(), [rowVectors.length, 1, rowVectors[0].length]);
    const output = await sess.run({'input_38': inputTensor});
    console.log(output);
}

runInference();

报错信息 #

2023-09-12 09:25:00.950381910 [E:onnxruntime:, sequential_executor.cc:514 ExecuteKernel] 
Non-zero status code returned while running Squeeze node. 
Name:'model_46/lambda_41/strided_slice__288' Status Message: 
/onnxruntime_src/onnxruntime/core/providers/cpu/tensor/squeeze.h:52 static
onnxruntime::TensorShapeVector onnxruntime::SqueezeBase::ComputeOutputShape(const 
onnxruntime::TensorShape&, const TensorShapeVector&) input_shape[i] == 1 was false. 
Dimension of input 1 must be 1 instead of 0. shape={1,0,18}

Uncaught:
Error: Non-zero status code returned while running Squeeze node. 
Name:'model_46/lambda_41/strided_slice__288' Status Message: 
/onnxruntime_src/onnxruntime/core/providers/cpu/tensor/squeeze.h:52 static 
onnxruntime::TensorShapeVector onnxruntime::SqueezeBase::ComputeOutputShape(const 
onnxruntime::TensorShape&, const TensorShapeVector&) input_shape[i] == 1 was false. 
Dimension of input 1 must be 1 instead of 0. shape={1,0,18}

问题原因 #

1. 输入维度不匹配：模型输入定义为Input(shape=(None, 16))（3D张量：[batch_size, seq_len, 16]），但训练时传入的是2D数组fd（[100,16]），TensorFlow自动扩展为[100,1,16]，但转ONNX时保留了动态维度的不确定性。
2. 切片逻辑错误：自定义Lambda层中的x[:,4]是对序列维度（第二个维度）取索引4的元素，当实际输入的序列长度为1时，该索引超出范围，导致生成维度为0的张量，触发ONNX Runtime的Squeeze节点错误（Squeeze要求被压缩的维度必须为1，不能是0）。
3. 输出形状声明错误：手动指定output_shape=(1,)与实际张量形状不匹配，干扰了TensorFlow到ONNX的形状推断。

解决方案 #

方案1：调整输入维度为固定2D（适合无序列场景） #

如果数据没有序列维度，直接将模型输入改为2D：

# 修改输入定义
I = Input(shape=(16,))  # 2D张量：[batch_size, 16]

# 自定义层修正为特征维度切片
def custom_layer(x):
    # x现在是2D张量：[batch_size, 18]
    return tf.exp(-x[:, 4] / (1e-9 + x[:, 0]) - x[:, 5] / (1e-9 + x[:, 1]))

# 移除output_shape参数，让TensorFlow自动推断
hl = tf.keras.layers.Lambda(custom_layer)(c)

同时调整ONNX Runtime输入张量维度：

const inputTensor = new ort.Tensor('float32', rowVectors.flat(), [rowVectors.length, rowVectors[0].length]);

方案2：适配3D输入的切片逻辑（保留序列维度） #

如果确实需要处理序列数据，修正切片逻辑为针对特征维度（最后一维）：

def custom_layer(x):
    # x是3D张量：[batch_size, seq_len, 18]
    # 取特征维度的第0、1、4、5位
    x0 = x[..., 0]  # [batch_size, seq_len]
    x1 = x[..., 1]
    x4 = x[..., 4]
    x5 = x[..., 5]
    result = tf.exp(-x4 / (1e-9 + x0) - x5 / (1e-9 + x1))
    # 扩展为3D张量，保持维度一致
    return tf.expand_dims(result, axis=-1)  # [batch_size, seq_len, 1]

hl = tf.keras.layers.Lambda(custom_layer)(c)  # 移除output_shape

转ONNX时显式声明动态维度：

tf2onnx.convert.from_keras(
    model, 
    output_file='tf2_neural_model_2deep.onnx',
    input_signature=[tf.TensorSpec((None, None, 16), tf.float32, name="input")]
)

方案3：避免动态维度的不确定性 #

转ONNX时固定输入维度，匹配训练时的实际输入形状：

tf2onnx.convert.from_keras(
    model, 
    output_file='tf2_neural_model_2deep.onnx',
    input_signature=[tf.TensorSpec((None, 1, 16), tf.float32, name="input")]
)

内容的提问来源于stack exchange，提问作者Mark C.