在C语言中运行输入为(-1,2)的TensorFlow SavedModel时出现段错误

我能正常在C语言中运行输入张量为(-1,1)的简易TensorFlow SavedModel，但将模型改为输入张量为(-1,2)时，出现段错误（Segmentation fault）。

使用版本 #

• gcc (x86_64-posix-seh-rev0, Built by MinGW-W64 project) 8.1.0
• TensorFlow 2.15.0（Python和C库均为此版本）

Python模型代码 #

import tensorflow as tf
import numpy as np

print(tf.__version__)

np.random.seed(0)
X = np.random.rand(100, 2)
y = 2 * X + 1 + 0.1 * np.random.randn(100, 1)

# 定义模型结构
model = tf.keras.Sequential([
    tf.keras.layers.Dense(1, input_shape=(2,))
])

# 编译并训练模型
model.compile(optimizer='sgd', loss='mean_squared_error')
model.fit(X, y, epochs=100)

tf.saved_model.save(model, 'simple_example')

模型检查输出 #

saved_model_cli show --dir simple_example --tag_set serve --signature_def serving_default

The given SavedModel SignatureDef contains the following input(s):
  inputs['dense_input'] tensor_info:
      dtype: DT_FLOAT
      shape: (-1, 2)
      name: serving_default_dense_input:0
The given SavedModel SignatureDef contains the following output(s):
  outputs['dense'] tensor_info:
      dtype: DT_FLOAT
      shape: (-1, 1)
      name: StatefulPartitionedCall:0
Method name is: tensorflow/serving/predict

C语言运行代码 #

#include <stdio.h>
#include <tensorflow/c/c_api.h>

int main(int argc, char *argv[])
{
  TF_Status* status = TF_NewStatus();
  TF_Graph*  graph  = TF_NewGraph();
  TF_Buffer* r_opts = TF_NewBufferFromString("",0);
  TF_Buffer* meta_g = TF_NewBuffer();
  
  const char* dir_name = "..\\data\\simple_example";
  const char* InputOperationName = "serving_default_dense_input";
  int64_t in_dims[] = {2,1};
  int64_t out_dims[] = {1,1};
  
  // 加载模型
  TF_SessionOptions* opts = TF_NewSessionOptions();
  const char* tags[] = {"serve"};
  TF_Session* session = TF_LoadSessionFromSavedModel(opts, r_opts, dir_name, tags, 1, graph, meta_g, status);
  if (session == NULL || TF_GetCode(status) != TF_OK )
    return -1;
  
  TF_Output input_op = {TF_GraphOperationByName(graph, InputOperationName ), 0};
  if (input_op.oper == NULL)
    return -1;

  // 定义输入张量
  TF_Tensor* input_tensor = TF_AllocateTensor(TF_FLOAT, in_dims, 2, in_dims[0] * in_dims[1] * sizeof(float));
  float* input_values = (float*)TF_TensorData(input_tensor);
  for (int n=0;n<(in_dims[0] * in_dims[1]);n++)
    input_values[n] = n;

  // 定义输出张量
  TF_Output output_op = {TF_GraphOperationByName(graph, "StatefulPartitionedCall"), 0};
  if (output_op.oper == NULL)
    return -1;
  TF_Tensor* output_tensor = TF_AllocateTensor(TF_FLOAT, out_dims, 2, out_dims[0] * out_dims[1] * sizeof(float));
  
  // 运行会话
  TF_SessionRun(session, NULL, &input_op, &input_tensor, in_dims[0]*in_dims[1], &output_op, &output_tensor, out_dims[0]*out_dims[1], NULL, 0, NULL, status);
  if (TF_GetCode(status) != TF_OK)
  {
      printf("failed: %s\n", TF_Message(status));
      return -1;
  }
  
  float* output_values = (float*)TF_TensorData(output_tensor);
  printf("Output value: %d\n", output_values[0]);
  
  return 0;
}

GDB调试输出 #

2024-06-21 14:10:16.608082: I tensorflow/cc/saved_model/loader.cc:217] Running initialization op on SavedModel bundle at path: ..\data\simple_example
2024-06-21 14:10:17.113501: I tensorflow/cc/saved_model/loader.cc:316] SavedModel load for tags { serve }; Status: success: OK. Took 3129450 microseconds.

Thread 1 received signal SIGSEGV, Segmentation fault.
0x00007ff980b2c346 in tensorflow!?CopyFromInternal@Tensor@tensorflow@@AEAAXAEBV12@AEBVTensorShape@2@@Z ()
   from tensorflow\tensorflow.dll

段错误的根源是输入张量维度不匹配和TF_SessionRun参数错误，修复点如下：

1. 修正输入维度
   模型期望输入形状为(-1,2)（任意批量大小，每个样本2个特征），但原代码中in_dims[] = {2,1}表示2个样本、每个样本1个特征，与模型要求完全相反。应改为int64_t in_dims[] = {1,2}（1个样本、2个特征），或{N,2}（N为自定义批量大小）。

2. 修正TF_SessionRun参数
   TF_SessionRun的第6个参数是输入张量的数量，第9个参数是输出张量的数量。原代码传入的是张量元素总数，实际输入和输出都只有1个张量，应都传入1。

3. 修正输出打印格式
   输出值是float类型，原代码用%d格式化整数会导致错误，改为%f即可。

修复后的关键代码片段：

// 修正输入维度：1个样本，2个特征
int64_t in_dims[] = {1,2};
int64_t out_dims[] = {1,1};

// ...

// 修正会话运行时的输入输出张量数量
TF_SessionRun(session, NULL, &input_op, &input_tensor, 1, &output_op, &output_tensor, 1, NULL, 0, NULL, status);

// ...

// 修正浮点值打印格式
printf("Output value: %f\n", output_values[0]);

另外，建议在程序结束时释放TensorFlow资源（如TF_DeleteTensor、TF_DeleteSession、TF_DeleteGraph等），避免内存泄漏，但这不是导致段错误的直接原因。

内容的提问来源于stack exchange，提问作者Martin